ArXiv: 2402.07598
🎯 Pitch
Surprisingly, learning full return distributions in reinforcement learning requires essentially no more samples than estimating mean returns alone, at least with a generative model. The authors prove this by introducing a new categorical algorithm that matches a recently established lower bound up to logarithmic factors, resolving a foundational open question. The key insight is a novel stochastic categorical CDF Bellman equation that cleanly handles the propagation of distributional uncertainty.
1. Executive Summary
This paper proposes a new model-based algorithm, the direct categorical fixed-point algorithm (DCFP), for distributional reinforcement learning and proves it achieves near-minimax-optimal sample complexity in the generative model regime—the first result of this kind for any distributional RL algorithm. The analysis operates on Markov reward processes with a generative model, estimating full return distributions rather than just expectations, using a categorical representation of probability distributions over a finite grid of outcomes together with a new stochastic categorical CDF Bellman equation (encoding the statistical fluctuations of repeated categorical dynamic programming updates with randomly sampled transitions). The central result establishes that DCFP requires only N = Ω(ε⁻²(1−γ)⁻³ polylog(|𝒳|/δ)) samples to achieve ε-accurate return distribution estimation in Wasserstein-1 distance with high probability, matching the minimax lower bound up to logarithmic factors and establishing that estimating full return distributions is essentially no more difficult, statistically speaking, than learning a value function—but only under the generative model access regime where independent samples can be drawn from every state.
2. Context and Motivation
The Core Problem: How Many Samples Does Distributional RL Actually Need?
The fundamental question this paper tackles can be stated concisely: if you want to learn full return distributions in a Markov decision process—rather than just expected returns—how many sampled transitions do you need? And critically, is this task statistically harder than simply estimating mean returns (value functions)?
This gap in understanding matters because distributional RL has become a core component of modern reinforcement learning systems, yet its statistical foundations remain comparatively unexplored. While the sample complexity of mean-return estimation has been thoroughly characterized—Azar et al. [2013] established a minimax lower bound of N = Ω(ε⁻²(1−γ)⁻³ log(|𝒳|/δ)) samples for ε-accurate value function estimation with high probability in the generative model regime, and provided algorithms (certainty equivalence) that achieve this bound—the analogous question for return distributions was unresolved. This is not merely a theoretical curiosity: understanding the sample complexity of distributional RL has direct practical implications for:
- Deployment efficiency: If distributional RL requires substantially more data than mean-return estimation, practitioners must weigh whether the richer signal (capturing aleatoric uncertainty, tail risk, and enabling risk-sensitive decision-making) justifies the increased data collection cost.
- Risk-sensitive applications: In domains like healthcare [Böck et al., 2022], robotics [Bodnar et al., 2020], and safety-critical systems, the full return distribution—not just its mean—directly informs decisions (e.g., via conditional-value-at-risk optimization). Knowing the statistical cost of acquiring these distributions is essential for designing reliable systems.
- Algorithm design: Without tight sample complexity bounds, it's unclear whether observed gaps between distributional RL algorithms reflect fundamental statistical hardness or suboptimal algorithm design. Upper bounds that are loose relative to lower bounds leave open the possibility that substantially better algorithms exist.
Conflicting Prior Evidence: A Lower Bound Without a Matching Upper Bound
Zhang et al. [2023] made significant progress on this question by proposing a distributional dynamic programming algorithm that operates with the empirical distributional Bellman operator ̂𝒯 (built from empirical transition probabilities). Their analysis produced two results with a notable gap:
Upper bound (what their algorithm achieves): N = Ω̃(ε⁻²(1−γ)⁻⁴ polylog(1/δ)) samples to guarantee ε-accurate estimation in Wasserstein-1 distance with high probability. The dependence on (1−γ)⁻¹ scales as a fourth power.
Lower bound (what is information-theoretically necessary): N = Ω̃(ε⁻²(1−γ)⁻³) samples, established by reduction to the mean-return case of Azar et al. [2013]. This follows because the mean is a functional of the full distribution, so any error in distribution estimation propagates to mean estimation—imposing at least the mean-return lower bound.
The gap between these bounds—Ω̃((1−γ)⁻⁴) in the upper bound versus Ω̃((1−γ)⁻³) in the lower bound—is substantial at typical discount factors. For γ = 0.99, a fourth-power dependence is roughly 100× larger than a third-power dependence, so this gap translates directly to orders-of-magnitude differences in practical sample requirements. Zhang et al. [2023] conjectured that their analysis was loose and that this gap could be closed, but left the question open.
The Practicality Gap: Unprojected Distributional DP Is Not Implementable
Beyond the statistical looseness, Zhang et al.'s [2023] algorithm suffers from a critical practical limitation that motivates much of this paper's technical development. Their approach applies the empirical distributional Bellman operator ̂𝒯 repeatedly without projecting the resulting distributions onto a finite representation. The problem is that each application of the distributional Bellman operator expands the support of the estimated return distributions:
- Start with return distributions supported on a finite set of outcomes.
- Apply
̂𝒯: the bootstrap function r(x) + γZ shifts and scales these outcomes, creating new outcome points that generally do not lie on the original support grid. - Without projection, these distributions grow in support complexity, quickly becoming computationally intractable.
Zhang et al. [2023] themselves acknowledge this: they had to introduce approximations to make their algorithm empirically runnable. This creates a tension between their theoretical analysis (which assumes no projection) and what is computationally feasible, leaving open the question of whether a practical algorithm can achieve the lower bound.
Where Categorical Dynamic Programming Falls Short as a Foundation
The natural candidate for a practical distributional RL algorithm is categorical dynamic programming (CDP), introduced by Bellemare et al. [2017] and analyzed by Rowland et al. [2018]. CDP represents return distributions as categorical distributions over a finite, equally-spaced grid of outcomes {z₁, …, zₘ} and applies a projection operator Πₘ after each distributional Bellman update to keep distributions supported on the grid. This tames the computational complexity: each update maps categorical distributions back to categorical distributions over the same grid.
Proposition 2.2 summarizes the key guarantee of Rowland et al. [2018]:
The operator Πₘ𝒯 is a contraction mapping with respect to the Cramér distance ℓ₂, with contraction factor √γ, and has a unique fixed point η_C ∈ 𝒫({z₁, …, zₘ})^𝒳. Further, the distance between η_C and the true return-distribution function η* can be bounded as:
This establishes CDP as principled—you can make the categorical approximation arbitrarily accurate by increasing m—but raises two questions critical to the current paper:
Question 1: How many categories m do you need? The bound above shows that to achieve ℓ₂-error ≤ ε, you need m ≥ 4(1−γ)⁻²ε⁻² + 1 categories. For tight guarantees at high discount factors, this grows quickly: at γ = 0.99 and ε = 0.01, you would need m ≈ 4 × 10⁸ categories—computationally prohibitive. However, this bound is derived for the representation error (how far the categorical fixed point is from the true return distribution), separate from the statistical error (how well you estimate the categorical fixed point from samples). The interaction between these two sources of error forms a central tension in the paper's analysis.
Question 2: Is CDP's iterative nature a bottleneck? The convergence rate √γ means that CDP requires O((1−γ)⁻¹ log(1/ε)) iterations to converge. At high discounts, this means many iterations—each requiring passes over the state space—even after acquiring the samples. This motivated the development of an alternative to iterative CDP.
The Iterative vs. Direct Solution Tension
CDP applies repeated iterations of η ← Πₘ𝒯η to converge to the categorical fixed point. But the contraction is slow when γ is close to 1 (contraction factor √γ ≈ 1), and each iteration incurs computational cost proportional to the state space and category count. For model-based approaches where you have already paid the statistical cost of acquiring samples, you might prefer a direct method that computes the categorical fixed point in one shot—solving a linear system rather than iterating to convergence. However, formulating CDP as a linear system with a unique solution is non-trivial because:
- The fixed-point equation F = T_P F (Equation 10, where F encodes cumulative distribution function values over the categorical grid) is not sufficient alone—the zero vector also satisfies it since distribution masses at each state are unconstrained.
- You need to explicitly enforce that F represents valid probability distributions (non-negative increments, total mass 1) while retaining linearity of the system.
- The natural representation is in cumulative rather than probability-mass form to make these constraints axis-aligned and tractable.
The paper's development of the DCFP algorithm—solving a carefully constructed linear system in ℝ^{𝒳×[m−1]}—resolves this tension, but its primary motivation is to enable the subsequent statistical analysis, not just computational convenience.
The Missing Statistical Framework for Categorical Distributional RL
Perhaps the deepest gap motivating this paper is the lack of appropriate analytical tools for studying statistical fluctuations in categorical distributional RL. Existing distributional RL theory focuses on:
- Convergence rates and approximation bounds for projection operators [Rowland et al., 2018]—these characterize how well the representation approximates the true return distributions, but say nothing about statistical errors from finite samples.
- Sample complexity of unprojected methods [Zhang et al., 2023]—these analyze what happens when you don't project, but are disconnected from practical projected algorithms.
What is missing is a framework that simultaneously handles:
- Representation error from using finitely many categories (controlled by m)
- Statistical error from using finitely many samples (controlled by N)
- The interaction between these two sources of error—does increasing m help statistically (by reducing representation bias) or hurt (by introducing more parameters to estimate)?
The paper's introduction of the stochastic categorical CDF Bellman equation (Section 5.2, developed fully in Appendix D) is designed precisely to fill this analytical gap. It encodes how random fluctuations from single-sample Bellman updates propagate through the categorical projection, and is used to bound the local-to-global amplification of estimation errors via a variance-like Bellman equation. Without this tool, the optimal (1−γ)⁻³ sample complexity cannot be established for the categorical approach.
The Generative Model Assumption: A Crucial Contextual Constraint
The paper's results hold under the generative model assumption [Kearns et al., 2002, Kakade, 2003], where the learning algorithm can independently sample next-state transitions from any state—effectively, a "reset and query" capability. This is a stronger access model than online RL (where the agent must navigate between states) but weaker than knowing the true transition probabilities. The generative model regime serves as a natural starting point for theoretical analysis of model-based algorithms because:
- It isolates the statistical difficulty (estimating transition probabilities from samples) from the exploration difficulty (navigating to states to collect samples).
- It allows direct comparison with the established mean-return results of Azar et al. [2013], which operate in the same regime.
- It is practically relevant for settings where simulation or offline data collection permits queries from arbitrary states.
A key point the paper does not address: whether similar optimality holds in the online RL setting, where exploration introduces additional dependencies between samples and typically inflates sample complexity by additional factors of (1−γ)⁻¹.
How This Paper Positions Itself
The paper defines its contribution not as proposing a new representation or a fundamentally different approach to distributional RL, but rather as closing a statistical gap while also providing a rigorous framework connecting categorical representations to sample complexity. Specifically:
-
It resolves the conjecture of Zhang et al. [2023]: there exists a practical, implementable algorithm (DCFP with m ≥ 4(1−γ)⁻²ε⁻² + 1 categories) that achieves the minimax lower bound up to logarithmic factors. The algorithm is not conceptually novel (it's model-based certainty equivalence with categorical projection), but proving it achieves the lower bound requires the new analytical machinery.
-
It establishes that distributional RL is not statistically harder than mean estimation: in the generative model regime, the sample complexity of learning ε-accurate return distributions in Wasserstein-1 matches that of learning ε-accurate value functions in L∞ norm (the Azar et al. [2013] bound). This is, as the paper notes, "perhaps surprising"—the full distribution is a much richer object than its expectation, yet requires asymptotically the same number of samples to estimate.
-
It provides practitioners with an analysis of practical tradeoffs: the empirical evaluation (Section 6) compares DCFP against existing distributional RL algorithms (CDP, QDP) across environments, discount factors, and atom counts, identifying regimes where each approach excels—notably finding that QDP benefits from its ability to adapt quantile locations in low-stochasticity settings, while DCFP excels under high discounts due to avoiding slow DP iterations.
The paper explicitly builds on:
- Rowland et al. [2018] for the categorical dynamic programming theory (projection, contraction, approximation bounds)
- Azar et al. [2013] and Agarwal et al. [2020] for the structure of the sample complexity proof (propagation of local errors, Bernstein concentration, analysis of empirical variance)
- Zhang et al. [2023] for the minimax lower bound and the problem framing
The novel contributions are the DCFP algorithm (Section 4), the stochastic categorical CDF Bellman equation (Section 5.2), and the synthesis of these components into a tight sample complexity proof (Theorem 5.1). The paper positions its analysis as primarily theoretical—the algorithm itself is the natural "certainty equivalent" approach for categorical distributional RL—but the proof structure reveals why it achieves optimality when properly configured.
3. Technical Approach
3.1 Reader Orientation
This paper develops a model-based distributional RL algorithm that estimates full return distributions (not just their means) from sampled transitions. The algorithm, called DCFP (direct categorical fixed-point), works by: (1) using samples to build an empirical model of the environment's transition probabilities, (2) representing return distributions as categorical distributions over a finite grid of possible outcomes, and (3) solving a carefully constructed linear system that directly computes the categorical fixed point—the distributional analog of "solve the Bellman equation" in value-based RL. The shape of the solution is: a generative model provides N i.i.d. next-state samples from each state, the algorithm forms empirical transition probabilities, selects a categorical grid with m support points (m controlled by desired accuracy ε and discount γ), builds a sparse matrix encoding the projected distributional Bellman operator in cumulative distribution function (CDF) form, strips out redundant coordinates, solves the resulting inhomogeneous linear system, and decodes the solution back to probability distributions—all while provably achieving the minimax-optimal sample complexity N = Ω̃(ε⁻²(1−γ)⁻³) up to logarithmic factors.
3.2 Big-Picture Architecture (Diagram in Words)
The system has five major components, each operating at a different stage of the pipeline:
-
Data collection (generative model): For each state x ∈ 𝒳, draw N i.i.d. next-state samples from the true transition kernel P(·|x). These samples are the sole statistical input to the algorithm—no online interaction, no trajectory data.
-
Empirical model construction: Build the empirical transition matrix P̂, where P̂(y|x) = N⁻¹ × (count of transitions from x to y). This is the classical "certainty equivalence" step—treat the sample frequencies as if they were the true probabilities.
-
Categorical grid definition: Choose m equally-spaced support points {z₁, …, zₘ} spanning the return interval [0, (1−γ)⁻¹] (or an environment-specific sub-interval if known). Return distributions will be approximated as categorical distributions over this grid. The number of categories m is chosen based on the target accuracy ε and discount γ: m ≥ 4(1−γ)⁻²ε⁻² + 1.
-
CDF operator construction: Build a matrix T_P̂ that encodes the projected distributional Bellman operator Πₘ𝒯 (with the empirical transition kernel P̂) acting on cumulative distribution function values rather than probability masses. This matrix has dimensions (|𝒳| × m) × (|𝒳| × m) but is highly sparse—at most 2⌈2/γ⌉ non-zero elements per row—enabling efficient sparse linear algebra.
-
Linear system solver (DCFP): Strip the m-th CDF coordinate (representing total mass, always 1) from each state's representation, converting the homogeneous system F = T_P̂ F into an inhomogeneous system (I − T̃_P̂) F̃ = H̃ over ℝ^{𝒳×[m−1]}. Solve this system directly using a sparse linear solver. The solution encodes the exact categorical fixed point—the distributional analog of the value function—in one shot, without requiring iterative dynamic programming.
Information flow: samples → empirical P̂ → categorical grid choice → matrix T_P̂ (using hat-function-based CDF coefficients) → reduced system (I − T̃_P̂)F̃ = H̃ → sparse linear solve → decoded return distribution estimates η̂(x) = Σᵢ (F̂ᵢ(x) − F̂ᵢ₋₁(x)) δ_{zᵢ}.
3.3 Roadmap for the Deep Dive
I'll explain the components in the order that reveals the logical dependencies:
- The categorical representation and its CDF encoding — because every downstream operation depends on understanding how return distributions are represented and why switching from probability masses to CDFs is the key enabling move.
- The projected distributional Bellman operator in CDF form — the core mathematical object that maps current categorical return-distribution estimates to updated ones, expressed as a linear operator T_P acting on CDF vectors.
- The DCFP linear system and its unique solution — how the fixed-point condition F = T_P F is converted into a solvable inhomogeneous system by exploiting the constraint that distributions sum to 1 and removing the redundant final CDF coordinate.
- The full model-based DCFP algorithm — assembling the pieces: from N samples per state to empirical P̂ to constructing the operator T_P̂ to solving the linear system to decoding the estimated return distributions.
- Design choices and their justifications — why CDF rather than probability-mass representation; why direct solve rather than iterative DP; why the generative model regime; why Wasserstein-1 as the target metric.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily a theoretical paper that develops an algorithm (DCFP) specifically to close a statistical gap—the algorithm is the natural "certainty equivalent" approach for categorical distributional RL, but its analysis requires substantial new machinery. The core idea is: if you represent return distributions as categorical CDFs over a finite grid and apply the empirical distributional Bellman operator as a linear transformation, you can directly solve for the fixed point and prove the resulting estimator achieves minimax-optimal sample complexity.
Categorical Representation and the Shift to CDF Encoding
In categorical distributional RL, a return distribution at state x is approximated as a categorical distribution over m equally-spaced outcomes z₁ < … < zₘ spanning [0, (1−γ)⁻¹]:
where δ_{z_i} is a Dirac (point mass) at z_i, p_i(x) ≥ 0 are the probability masses, Σᵢ p_i(x) = 1, and z_i = (i−1)/(m−1) × (1−γ)⁻¹ for i = 1, …, m.
What this represents: At each state x, the algorithm maintains a histogram over m equally-spaced return values. The parameter vector (p₁(x), …, pₘ(x)) must live in the probability simplex—each entry non-negative, sum to 1. The grid spacing is uniform with width Δ = (m−1)⁻¹(1−γ)⁻¹.
Why this representation: Categorical distributions form a finite-dimensional parameterization of probability distributions. The key tradeoff is that m controls expressivity—larger m means finer granularity (closer to the true continuous return distribution, as quantified by the ℓ₂ approximation bound from Rowland et al. [2018]), but at the cost of more parameters per state. This parameterization is special because the projected distributional Bellman operator Πₘ becomes a linear operator on the probability masses—a property exploited throughout the paper.
The CDF encoding decision. Rather than working directly with probability masses p_i(x), the paper reparameterizes the categorical distribution in terms of cumulative distribution function (CDF) values:
for i ∈ {1, …, m}. By convention, F₀(x) = 0 and Fₘ(x) = 1 always.
What this representation provides: F(x) ∈ ℝᵐ is a vector with F₁(x) ≤ F₂(x) ≤ … ≤ Fₘ(x) = 1, where each component F_i(x) is the cumulative probability mass up to and including atom z_i. The probability mass at atom i is recovered as p_i(x) = F_i(x) − F_{i−1}(x).
Why CDFs over probability masses: The constraint that the distribution sums to 1 becomes a simple axis-aligned linear constraint (Fₘ(x) = 1) rather than a sum constraint over all atoms (Σᵢ p_i(x) = 1). This is the critical technical enabler: when we later formulate the fixed-point condition as a linear system, the CDF representation makes the constraints on valid probability distributions take the form of fixing specific coordinates to known constants—something that allows us to strip those coordinates out and convert a homogeneous system with non-unique solutions into an inhomogeneous system with a unique solution (Proposition 4.2). In probability-mass space, the constraint is a hyperplane equation coupling all coordinates, which is harder to handle algebraically.
Monotonicity is preserved. The CDF representation naturally encodes monotonicity: F_{i}(x) ≤ F_{i+1}(x) by definition (cumulative mass is non-decreasing). When the operator T_P is applied, the paper verifies that this monotonicity is preserved—if the input F satisfies 0 ≤ F₁ ≤ … ≤ Fₘ = 1, then the output T_P F also satisfies these constraints (this follows from the fact that T_P = Πₘ𝒯 in CDF coordinates, and Πₘ𝒯 maps valid distributions to valid distributions).
The Projected Distributional Bellman Operator in CDF Form
The core computational primitive in categorical distributional RL is applying the projected distributional Bellman operator Πₘ𝒯 to a collection of categorical return-distribution estimates. This operator: (1) for each state x, computes the push-forward distribution of the return under the bootstrap function (r(x) + γZ, where Z ~ η(y) for the next state y), then (2) projects the resulting distribution back onto the grid {z₁, …, zₘ} using the categorical projection Πₘ.
From probability masses to the update rule. For an input RDF η(x) = Σⱼ pⱼ(x) δ_{zⱼ} in probability-mass form, the probability masses pᵢ′(x) of (Πₘ𝒯 η)(x) are given by Equation (6) in the paper:
where h^x_{i,j} = h_i(r(x) + γzⱼ), and h_i is the "hat function" centered at grid point z_i.
What h_i computes: The hat function h_i : [0, (1−γ)⁻¹] → [0, 1] linearly interpolates between 1 at z_i and 0 at neighboring grid points z_{i−1} and z_{i+1}, and is zero outside [z_{i−1}, z_{i+1}]. For a value r(x) + γzⱼ (the shifted-and-scaled atom from the backed-up distribution), h_i(r(x) + γzⱼ) determines what fraction of the backed-up mass at outcome r(x) + γzⱼ should be allocated to grid point z_i—it interpolates between adjacent grid points in proportion to their proximity. This implements the categorical projection Πₘ: split each backed-up atom between its two nearest grid points.
Converting to CDF form. Rather than accumulating probability masses, the paper sums over i to get the cumulative relationship. Define the cumulative sum of hat functions up to index i:
for j = 1, …, m, with convention H^x_{i, m+1} = 0.
What H^x_{i,j} represents: This is the total fraction of the mass from backed-up atom zⱼ (shifted to r(x) + γzⱼ) that gets assigned to grid points z₁ through z_i—i.e., the cumulative mass at or below z_i from the projection of δ_{r(x)+γzⱼ}. Since Σ_{l=1}ᵐ h_l(z) = 1 for any z ∈ [0, (1−γ)⁻¹] (the hat functions form a partition of unity), we have H^x_{m,j} = 1 always.
Then Proposition 4.1 gives the CDF form of the update:
What this equation computes: For each state x and each cumulative index i (1 through m), the updated CDF value Fᵢ′(x) is a weighted sum over next states y and atom indices j. The weight (H^x_{i,j} − H^x_{i,j+1}) is the mass that atom j at the next state contributes to the cumulative mass up to z_i at state x after bootstrapping and projection. The product (H^x_{i,j} − H^x_{i,j+1}) × Fⱼ(y) integrates over the CDF of the next-state distribution. The outer sum P(y|x) mixes over transition probabilities.
Why the H^x_{i,j} − H^x_{i,j+1} difference appears: This is a telescoping sum artifact from converting probability masses to CDFs. In probability-mass form, the update involves pⱼ(y) = Fⱼ(y) − F_{j−1}(y). Substituting and rearranging the sums yields the difference-of-H form, which is algebraically cleaner for the matrix representation that follows.
The linear operator T_P. Equation (7) defines a linear transformation on the space ℝ^{𝒳×m} (CDF values for all state-index pairs). Writing F and F′ as vectors in this space, the update is:
where T_P is an |𝒳|m × |𝒳|m matrix with entries:
What T_P physically does: Apply the projected distributional Bellman operator Πₘ𝒯 to a collection of categorical return distributions represented by their CDF vectors, producing updated CDF vectors. The matrix is defined entirely from the transition probabilities P and the hat-function coefficients H^x_{i,j} (which themselves depend only on the rewards r(x) and the grid {z_i}—not on the current distribution estimates). This linearity is the key property that enables everything that follows.
Why this matrix form matters: Since T_P is linear, the fixed point of CDP satisfies F* = T_P F*, a linear system. This transforms the problem of finding the categorical fixed point from an iterative procedure (CDP) to a direct algebraic one (solving a linear system). The matrix T_P is highly structured: for each state x, the submatrix encoding the local backup-and-project operation can be written as P_x ⊗ B_x, where P_x is the x-th row of P and B_x ∈ ℝᵐˣᵐ captures the shifting/scaling/projecting at state x.
Sparsity of T_P. As detailed in Appendix G.3, for most (x, i; y, j) index pairs, H^x_{i,j} − H^x_{i,j+1} = 0. This non-zero condition requires that at least one of r(x) + γzⱼ or r(x) + γz_{j+1} lies in the interval [z_i, z_{i+1}]. For fixed x and i, at most ⌈2/γ⌉ indices j can satisfy this (since the distance between successive backed-up atoms is γ × Δ). Thus each row of T_P has at most 2⌈2/γ⌉ non-zero elements—constant in m and |𝒳| for γ not too small. This sparsity is crucial for practical implementations, enabling both the iterative CDP (sparse matrix-vector products) and DCFP (sparse linear system solution) to scale.
The DCFP Linear System and Its Unique Solution
The fixed-point condition F = T_P F (Equation 10) is a homogeneous linear system (I − T_P)F = 0. However, this system does not have a unique solution—the all-zero vector F = 0 is always a solution, as are multiples of the all-ones vector in the m-th coordinate (since Fₘ is unconstrained by the equation).
The problem with the homogeneous system. The homogeneous system (I − T_P)F = 0 encodes the condition that F is a fixed point of T_P, but it ignores the requirement that F represents valid probability distributions (Fₘ(x) = 1 for all x). Any vector in the nullspace of (I − T_P) is a fixed point, but only one corresponds to the actual categorical fixed point η_C. The contraction property (Proposition 2.2) guarantees uniqueness of η_C in the space of valid probability distributions, but the linear system alone doesn't know about this constraint.
Proposition 4.2: Converting to an inhomogeneous system. The key insight is to:
- Explicitly impose Fₘ(x) = 1 for all x ∈ 𝒳 (the total mass constraint).
- Remove the variables Fₘ(x) from the system (since they're now known constants).
- Remove the rows corresponding to i = m from the system (since those equations become redundant—they reduce to 1 = 1 after substituting Fₘ(x) = 1).
The result is a linear system in ℝ^{𝒳×[m−1]} (one fewer coordinate per state), which becomes inhomogeneous:
where:
- F̃ ∈ ℝ^{𝒳×[m−1]} is the vector of unknown CDF values for indices 1 through m−1 at each state
- T̃_P(x,i; y,j) = P(y|x)(H^x_{i,j} − H^x_{i,j+1}) for 1 ≤ i,j ≤ m−1 (the submatrix of T_P with the m-th row and column stripped)
- H̃(x,i) = H^x_{i,m} (the contribution from the fixed Fₘ(y) = 1 terms)
What this equation computes: For each state x and CDF index i < m, the left side (I − T̃_P)F̃ represents "Fᵢ(x) minus what the Bellman backup would predict for Fᵢ(x) if all distributions summed to 0 at the top." But because the true distributions DO sum to 1 at Fₘ, the right side H̃(x,i) = H^x_{i,m} captures the mass contributed by the guaranteed-1 tail—it's the cumulative probability up to z_i that comes from backing up the certain event that the backed-up distribution has total mass 1 below zₘ.
Why the inhomogeneous form has a unique solution. Proposition 4.3 establishes uniqueness: suppose F̃₁ and F̃₂ are two distinct solutions. Reconstructing the full CDF vectors F₁ and F₂ (by appending Fₘ(·) = 1) would give two distinct fixed points of Πₘ𝒯 in the space of valid probability distributions. But Proposition 2.2 (Rowland et al. [2018]) established that Πₘ𝒯 is a contraction on this space with respect to Cramér distance, so its fixed point is unique—contradiction. This is formalized in Appendix C using the norm ∥·∥_{ℓ₂,∞} on ℝ^{𝒳×m} to show T_P is a contraction on the subspace {F : Fₘ(x) = 0 for all x}.
The norm ∥·∥_{ℓ₂} on ℝᵐ. Throughout the analysis, the paper uses a rescaled ℓ₂ norm:
where F ∈ ℝᵐ is a CDF vector.
What this norm measures: Under this norm, if two categorical distributions (supported on {z₁, …, zₘ}) have CDF vectors F and F′, then ∥F − F′∥{ℓ₂} equals exactly their Cramér distance ℓ₂—the integral of the squared difference of their CDFs. The factor 1/(m(1−γ)) corrects for the discrete approximation: it's the quadrature weight for the integral ∫₀^{(1−γ)⁻¹} (F(t) − F′(t))² dt discretized at the m grid points. The supremum version ∥F∥{ℓ₂,∞} = max_x ∥F(x)∥_{ℓ₂} extends this to the state space, capturing the worst-case Cramér distance across states.
Contraction of T_P. Proposition B.1 establishes that T_P is a contraction on the subspace {F : Fₘ(x) = 0 for all x} with respect to ∥·∥{ℓ₂,∞}, with contraction factor √γ. This is the CDF-translated version of the contraction result from Rowland et al. [2018] (Proposition 2.2). The subspace condition Fₘ = 0 is crucial: it means we're looking at differences between valid CDF vectors (both have Fₘ = 1, so their difference has Fₘ = 0). This contraction property is what guarantees uniqueness—if F₁ and F₂ are two fixed points, their difference d = F₁ − F₂ satisfies d = T_P d and has dₘ = 0, so ∥d∥{ℓ₂,∞} = ∥T_P d∥{ℓ₂,∞} ≤ √γ ∥d∥{ℓ₂,∞}, forcing d = 0.
Solving the DCFP linear system in practice. Algorithm 1 summarizes DCFP:
- Compute the matrices Hˣ for each state x (via Equation 8, summing the hat functions).
- Build T̃_P and H̃ from the transition probabilities P and the Hˣ matrices (Equation 9).
- Solve (I − T̃_P)F̃ = H̃ using a sparse linear solver (e.g.,
scipy.sparse.linalg.spsolve). - The full CDF solution F* is obtained by appending Fₘ(x) = 1 to F̃ for each x.
Computational complexity. The matrix T̃_P has dimensions |𝒳|(m−1) × |𝒳|(m−1). Solving the linear system with a dense solver would cost O(|𝒳|³m³), which is prohibitive. However, the matrix is highly sparse (approximately 2⌈2/γ⌉ non-zeros per row, independent of m and |𝒳| for fixed γ), enabling sparse direct solvers to achieve much better practical performance. Further, for factored or structured state spaces, additional sparsity from P (which may have far fewer than |𝒳|² non-zero entries) further reduces the effective dimension.
Why direct solve rather than iterative CDP? For the theoretical analysis, DCFP is essential because it gives the exact empirical categorical fixed point F̂, enabling a clean separation between representation error (F* vs. η*) and statistical error (F̂ vs. F*). For practical computation, DCFP avoids the slow √γ-rate convergence of CDP, which requires O((1−γ)⁻¹ log(1/ε)) iterations to reach ε-accuracy—unacceptably slow when γ ≈ 1. The experiments (Section 6) confirm that DCFP substantially outperforms CDP in wallclock time at high discount factors. However, Appendix F.1 shows that CDP with sufficiently many iterations (k ≥ 2 log(1/ε) / log(1/γ)) also achieves the optimal sample complexity—the choice between DCFP and CDP is computational, not statistical.
Model-Based DCFP: From Samples to Return Distributions
The full model-based DCFP algorithm operates under the generative model access pattern: for each state x ∈ 𝒳, we observe N i.i.d. samples Xˣ₁, …, Xˣ_N from the true transition distribution P(·|x).
Step 1: Build empirical transition probabilities. The empirical transition matrix P̂ is constructed using simple frequency counting:
for each state pair (x, y). This is the standard nonparametric maximum-likelihood estimate of a categorical distribution from i.i.d. draws. No smoothing, regularization, or Bayesian prior is applied—the algorithm uses raw empirical frequencies.
Step 2: Choose the number of categories m. Based on the target accuracy ε and discount γ, set:
Why this specific threshold: This choice ensures that the representation error (ℓ₂ distance between the true categorical fixed point F* and the true return distribution η*) is bounded by ε/2 (via the Rowland et al. [2018] bound in Equation 5). The remaining ε/2 budget is allocated to the statistical error (ℓ₂ distance between the estimated F̂ and the true categorical fixed point F*). The factor 4 comes from the specific form of the Cramér-distance approximation bound, which scales as 1/((1−γ)√(m−1)). To achieve ≤ ε/2, we need 1/((1−γ)√(m−1)) ≤ ε/2, which rearranges to m−1 ≥ 4(1−γ)⁻²ε⁻².
Critical observation: The threshold for m does not depend on N. This means there is no statistical penalty for using a large number of categories—you can make m arbitrarily large (to reduce representation error) without hurting sample complexity, because the number of parameters being estimated (|𝒳| × m) does not appear in the sample complexity bound (only the number of states |𝒳| does). This is a surprising and important result: the categorical representation's statistical cost is controlled by the state space size, not the category count.
Step 3: Build the CDF operator with P̂. Construct the quantities Ĥˣ_{i,j} = Σ_{l≤i} h_l(r(x) + γzⱼ) using the known reward function r(x) and the chosen grid. These depend only on the reward function and grid geometry, not on the transition samples. Then build T_P̂ and H̃ using P̂ in place of P in the formulas of Proposition 4.2. The matrix T_P̂ inherits sparsity from both P̂ (if the true P is sparse) and from the atomic structure discussed earlier.
Step 4: Solve the empirical linear system. Solve (I − T̃_P̂)F̃ = H̃ to obtain the estimated CDF values F̃̂. Append F̂ₘ(x) = 1 for all x to obtain the full estimated CDFs F̂ ∈ ℝ^{𝒳×m}.
Step 5: Decode to return distributions. Convert from CDF values back to probability distributions:
with the convention F̂₀(x) = 0. Each F̂_i(x) − F̂_{i−1}(x) is non-negative because F̂ was constructed as a fixed point of T_P̂, which preserves the monotonicity constraints of a valid CDF.
What the output represents: For each state x, η̂(x) is a categorical distribution over {z₁, …, zₘ} that approximates the true return distribution η*(x). The quality of this approximation in Wasserstein-1 distance is bounded by ε with high probability 1−δ, provided N satisfies the sample complexity bound in Theorem 5.1.
The Wasserstein-1 metric as target. The paper measures error in Wasserstein-1 distance:
Why Wasserstein-1: This metric is particularly compatible with categorical distributional RL because: (1) it can be bounded above by Cramér distance (Lemma 5.2: w₁ ≤ (1−γ)^{−1/2} ℓ₂), which is the natural metric for the categorical contraction theory; (2) it upper-bounds differences in expectations of Lipschitz functions, making it a practically meaningful error measure for downstream decision-making; (3) it controls differences in risk measures like conditional value-at-risk (CVaR). The reduction w₁ ≤ (1−γ)^{−1/2} ℓ₂ costs a factor of (1−γ)^{−1/2} in the conversion, which is why the final sample complexity picks up an extra (1−γ)^{−1} factor relative to the ℓ₂ bound (going from Ω̃(ε⁻²(1−γ)⁻²) in ℓ₂ to Ω̃(ε⁻²(1−γ)⁻³) in w₁).
Design Choices and Their Justifications
Choice 1: Categorical representations over other distribution parameterizations. The paper focuses on categorical distributions (histograms over a fixed grid) rather than quantile representations [Dabney et al., 2018b] or particle-based methods [Morimura et al., 2010a] for several reasons: (a) The projected Bellman operator Πₘ𝒯 is linear in the probability masses (and hence in the CDF values), which enables the direct linear-system formulation of DCFP. Quantile dynamic programming (QDP), by contrast, involves a non-linear projection (sorting and re-quantiling), making it challenging to analyze as a linear system. (b) The existing convergence and approximation theory from Rowland et al. [2018] provides rigorous bounds on representation error (Equation 5) that are essential for the proof's decomposition. (c) The number of categories m provides an explicit, tunable knob for trading expressivity against computation, with the theoretical guarantee that increasing m never hurts sample complexity.
Choice 2: CDF encoding over probability-mass encoding. The shift to CDF values is not cosmetic—it is essential for formulating the fixed-point condition as a linear system with a unique solution. In probability-mass space, the constraint Σᵢ pᵢ(x) = 1 couples all m coordinates. Removing one coordinate (say pₘ) and substituting pₘ = 1 − Σ_{i<m} pᵢ works but introduces non-axis-aligned constraints and complicates the linear algebra. In CDF space, the constraint Fₘ(x) = 1 is axis-aligned—it involves a single coordinate per state. This allows the clean reduction from the homogeneous system (I − T_P)F = 0 to the inhomogeneous system (I − T̃_P)F̃ = H̃ by simply dropping the m-th row and column and moving the known Fₘ terms to the right-hand side.
Choice 3: Direct linear solve (DCFP) over iterative CDP. While Appendix F.1 shows that CDP with sufficiently many iterations also achieves the optimal sample complexity, DCFP is preferred for both theoretical and practical reasons: (a) DCFP computes the exact empirical fixed point, eliminating iteration error from the analysis and enabling a clean ℓ₂(F̂, F*) bound without an additional O(γ^k) term. (b) CDP converges at rate √γ, requiring O((1−γ)⁻¹ log(1/ε)) iterations—impractical when γ = 0.99 (√γ ≈ 0.995, needing ~4600 iterations for 10⁻⁴ error). DCFP solves the system in one shot, at the cost of a sparse linear solve whose complexity depends on the sparsity pattern, not the contraction rate. (c) The linear system formulation reveals structure (sparsity, contraction properties) that the proof exploits.
Choice 4: Generative model regime. The assumption of N i.i.d. samples per state isolates the statistical challenge of distribution estimation from the exploration challenge. This allows direct comparison with the foundational results of Azar et al. [2013] and Zhang et al. [2023] and makes the minimax lower bound tight (the lower bound is proven in this same regime, so matching it exactly is possible). Extending to online RL would require additional techniques to handle the dependencies between samples collected along trajectories and would likely introduce additional (1−γ)⁻¹ factors.
Choice 5: Wasserstein-1 as the evaluation metric. The paper uses w₁ as the primary metric because it is a standard, interpretable distance between probability distributions that controls differences in expectations of Lipschitz functions and risk measures like CVaR. The conversion from ℓ₂ (the natural metric for categorical contraction analysis) to w₁ via Lemma 5.2 introduces a (1−γ)^{−1/2} factor, which propagates to the final sample complexity. This explains why the optimal sample complexity in w₁ (ε⁻²(1−γ)⁻³) has one more power of (1−γ)⁻¹ than the ℓ₂ bound (ε⁻²(1−γ)⁻²): the metric conversion costs a factor when the return range is O((1−γ)⁻¹).
Choice 6: Global atom locations vs. environment-specific ones. The theoretical analysis uses atoms spanning the entire possible return interval [0, (1−γ)⁻¹] (the worst-case return range). This ensures the guarantee holds without any prior knowledge of the environment. However, the experiments (Appendix G) show that using environment-specific atom locations—e.g., [r_min(1−γ)⁻¹, r_max(1−γ)⁻¹] when the reward range is known to be [r_min, r_max]—can substantially improve practical accuracy for a given m, since the same number of atoms are now concentrated in a smaller interval, increasing effective resolution. This tighter localization is compatible with the theory (the proof only requires the grid to cover the actual support of return distributions) and is recommended for practitioners.
Choice 7: No dependence on m in sample complexity. Perhaps the most surprising and consequential design consequence of the analysis: Theorem 5.1 shows that N does not depend on m (beyond the condition that m is sufficiently large to control representation error). This means you can freely increase m to reduce approximation error without statistical penalty. The intuition is that the Cramér distance integrates over the return interval, so estimation errors at individual atoms "average out" in the integral—the effective number of degrees of freedom being estimated is controlled by the smoothness of the CDF (which is O(1)) rather than the number of discrete atoms (which is m). This is formalized through the Hilbert-space Bernstein inequality used in the proof, which avoids per-coordinate union bounds that would introduce log(m) factors.
4. Key Insights and Innovations
Innovation 1: Distributional RL Is Statistically No Harder Than Mean-Value Estimation
The paper's most fundamental contribution is not a new algorithm, but a provable equivalence: in the generative model regime, estimating full return distributions (an infinite-dimensional object) requires essentially the same number of samples as estimating just their expectations (a scalar per state). This resolves a conjecture by Zhang et al. [2023] and, more importantly, reframes what distributional RL "costs" statistically.
Before this work, the field faced a troubling gap. Zhang et al. [2023] established a lower bound of Ω̃(ε⁻²(1−γ)⁻³) samples for distribution estimation—matching the value-function lower bound of Azar et al. [2013]—but their best algorithm required Ω̃(ε⁻²(1−γ)⁻⁴), an extra factor of (1−γ)⁻¹ that could translate to orders-of-magnitude differences at high discount factors. Practitioners reading this literature might reasonably conclude that distributional RL carries a genuine statistical premium: you pay extra samples to get richer predictions. The gap was not obviously loose—distributional estimates ARE richer objects, so perhaps the extra cost was fundamental.
What the DCFP analysis proves is that this premium is an artifact of suboptimal algorithm design, not an information-theoretic necessity. Theorem 5.1 shows that with m = 4(1−γ)⁻²ε⁻² + 1 categories and N = Ω̃(ε⁻²(1−γ)⁻³) samples, DCFP achieves ε-accurate estimates in Wasserstein-1 distance. The sample complexity matches the lower bound up to logarithmic factors. This closes the gap and establishes that the statistical difficulty of distributional RL is fundamentally governed by the same (1−γ)⁻³ scaling as mean-value estimation—not (1−γ)⁻⁴ or worse.
The significance extends beyond this specific numeric match. The proof reveals that the categorical representation, when properly analyzed, does not introduce additional statistical degrees of freedom. Despite having m parameters per state, the effective complexity is controlled by the smoothness of the CDF (integrated over the return interval), not by the number of discrete atoms. The Hilbert-space Bernstein inequality (Appendix E.4) avoids per-coordinate union bounds that would introduce log(m) or worse dependencies—the Cramér distance integrates across atoms, so errors average out rather than accumulating. This is a fundamental insight about why distributional RL can be statistically efficient: the structure of the distributional Bellman equation, combined with an appropriate metric (Cramér/Wasserstein), provides implicit regularization that prevents the curse of dimensionality from the representation choice.
This result also has immediate practical implications. It means practitioners can invest in distributional RL for the downstream benefits (risk-sensitive decision-making, better exploration, richer representations) without worrying about a hidden statistical tax. There is no fundamental tradeoff between prediction richness and sample efficiency—at least in the generative model regime. The caveat, of course, is that this equivalence holds for the generative model setting; in online RL, exploration costs may introduce additional gaps that are not yet characterized.
Innovation 2: The Stochastic Categorical CDF Bellman Equation as a Diagnostic Tool for Distributional Uncertainty
The paper introduces the stochastic categorical CDF Bellman equation (SC-CDF Bellman equation, Definition 5.6), a new type of distributional Bellman equation that encodes the statistical fluctuations of categorical dynamic programming under random sampling. This is not an algorithm—it's a conceptual tool that bridges two previously disconnected perspectives on distributional RL.
To understand why this is innovative, consider the dual nature of distributional RL predictions. The return distribution η*(x) captures aleatoric uncertainty—the inherent randomness in the return due to stochastic transitions and rewards. But when we estimate η*(x) from finite samples, there is also epistemic uncertainty—our uncertainty about what the true return distribution actually is, due to having seen only N samples. Prior distributional RL theory addressed these completely separately: Bellemare et al. [2017] and Rowland et al. [2018] studied aleatoric uncertainty (how well categorical distributions approximate true return distributions), while Zhang et al. [2023] and Azar et al. [2013] studied epistemic uncertainty (how sample size affects estimation accuracy), but using fundamentally different mathematical objects—distributions over returns versus confidence intervals around estimates.
The SC-CDF Bellman equation unifies these: its solution Φ_Q(x) is a distribution over CDF vectors—a doubly-distributional object. The outer distribution (over CDFs) encodes epistemic uncertainty: each realization Φ_Q(x) is a valid CDF of a return distribution, and the spread of these realizations around the mean F_Q(x) = E[Φ_Q(x)] captures how much our estimate would vary under repeated sampling. The inner CDFs encode aleatoric uncertainty: each realization summarizes a possible return distribution. This is exactly the right object for analyzing sample complexity: the local squared-Cramér variation σ_Q (Definition 5.5) measures how much a single-sample Bellman update perturbs the fixed-point CDF, and the global squared-Cramér variation Σ_Q (Definition 5.9) measures the accumulated effect of repeated such perturbations.
What makes this genuinely novel—rather than a minor reformulation—is that it provides a variance Bellman equation for distributional RL. The inequality in Proposition 5.11:
is the distributional analog of the variance Bellman equation used by Azar et al. [2013] for mean-return sample complexity. In the mean-return case, the variance Bellman equation relates the variance of the value estimate to the variance of immediate rewards and the discounted variance at successor states, enabling a bound of the form ∥(I − γP)⁻¹σ∥∞ ≤ 2/(1−γ). The SC-CDF Bellman equation plays the same role for Cramér-distance-based variances, but with an additional correction term (the subtracted constants) that arises from the categorical projection's finite resolution—an artifact with no mean-return analog.
The broader significance is that this equation is likely reusable beyond this paper's specific analysis. Any future work studying statistical efficiency of categorical distributional RL—whether in online settings, with function approximation, or under different sampling assumptions—will need some way to propagate local estimation errors to global guarantees. The SC-CDF Bellman equation provides the template for doing so. It is, in essence, the missing analytical bridge between the deterministic contraction theory of Rowland et al. [2018] and the stochastic concentration arguments needed for sample complexity.
Innovation 3: The CDF Reparameterization Enables a Direct Linear-System Formulation of Categorical DP
While the DCFP algorithm itself is a natural "certainty equivalent" approach, the specific technical move that makes it work—reparameterizing categorical distributions from probability masses to cumulative distribution functions—is a conceptual shift with implications beyond this paper. It converts the projected distributional Bellman operator into a sparse linear system with a unique solution, enabling both the theoretical analysis and practical implementation.
Prior work on categorical distributional RL [Bellemare et al., 2017, Rowland et al., 2018] worked in probability-mass space: distributions are represented by vectors p ∈ ℝᵐ with p_i ≥ 0 and Σᵢ pᵢ = 1. The projected Bellman update maps p to p′ via Equation (6), which is linear in p but involves hat-function coefficients hˣ_{i,j} that are inconvenient for fixed-point analysis. The constraint Σᵢ pᵢ = 1 couples all m coordinates, making it difficult to isolate a linearly independent subset. In this representation, solving for the fixed point requires either iterating to convergence (CDP) or solving a constrained linear system with a Lagrange multiplier—both messy for theoretical analysis.
The CDF reparameterization F_i = Σ_{j≤i} pⱼ transforms these constraints into simple boundary conditions: Fₘ = 1 always, and monotonicity is F_i ≤ F_{i+1}. The critical advantage is that Fₘ = 1 is axis-aligned—it involves exactly one coordinate per state, not a sum over all coordinates. This allows the clean reduction in Proposition 4.2: strip out the m-th coordinate (now a known constant), remove the m-th equation (now redundant), and solve an inhomogeneous system (I − T̃_P)F̃ = H̃ over ℝ^{𝒳×[m−1]}. The matrix T̃_P is sparse (≈ 2⌈2/γ⌉ non-zeros per row independent of m and |𝒳|), and the system has a unique solution by contraction arguments.
The significance of this move extends beyond algorithmic convenience. The CDF representation naturally interfaces with the Cramér distance ℓ₂—the metric in which Πₘ𝒯 contracts and in which the approximation bounds of Rowland et al. [2018] are expressed. The ℓ₂ norm on CDF vectors (∥F∥²_{ℓ₂} = (m(1−γ))⁻¹ Σᵢ F²ᵢ) is precisely the quadrature approximation to the Cramér integral ∫₀^{(1−γ)⁻¹} F²(t) dt, making the CDF representation the "native" coordinates for the contraction theory. This alignment between representation, metric, and contraction is not accidental—it is what enables the tight sample complexity analysis to work without losing factors from metric mismatches.
For practitioners, this reframing suggests that CDF-valued dynamic programming may be preferable to probability-mass-valued DP even when the final output is desired in probability-mass form. The linear structure, sparsity, and direct solvability of the CDF formulation are advantages that extend to any tabular categorical distributional RL implementation, not just the generative model setting.
Innovation 4: Verifier Over-Optimization at the Algorithmic Level—Why DCFP Wins Where QDP Loses
The paper's experimental findings (Section 6) provide a diagnostic insight that is not captured by the worst-case sample complexity theory but is practically significant: the relative performance of categorical (DCFP) versus quantile (QDP) methods depends critically on environment structure—specifically, the presence of short high-probability self-loops and the level of stochasticity.
In the two-state environment with γ = 0.99 (Figure 9), QDP with 30 atoms achieves substantially worse Wasserstein-1 error than DCFP with the same atom count—a gap that persists even as both methods improve with more atoms. This is not a sample-size issue (N = 10⁶ per state is large) but a fixed-point approximation error issue: QDP's quantile representation, which adapts its support points to concentrate mass where the distribution is concentrated, can develop systematic biases in environments where the return distribution has sharp features that are poorly aligned with the quantile support. Rowland et al. [2024] previously identified this phenomenon in the context of quantile temporal-difference learning; the current paper's experiments extend this diagnostic to the model-based setting, showing it is a property of the representation, not the learning algorithm.
Conversely, QDP excels in the "low random" environment (Figure 7) where transition probabilities are near-deterministic (Dirichlet(0.01, …, 0.01) concentrates mass on few transitions). Here, QDP's ability to concentrate its quantile support on the few relevant return values gives it a representation-efficiency advantage over DCFP's uniform grid, which must spread atoms across the full return interval even if most of the probability mass occupies a narrow region.
The practical insight is that there is no universal winner between categorical and quantile distributional RL—the choice should be informed by environment diagnostics. For near-deterministic environments with concentrated return distributions, quantile methods exploit their adaptive support to achieve better accuracy per parameter. For environments with diffuse return distributions or sharp bootstrap-induced features (as in tight self-loops), categorical methods' uniform coverage and linear structure provide robustness. The paper's experiments also reveal that DCFP's direct linear solve substantially outperforms CDP's iterative approach in wallclock time at high discounts (γ = 0.99), where CDP requires thousands of iterations for convergence—a pragmatic finding that makes DCFP attractive even when the theoretical optimality is not the primary concern.
This diagnostic dimension—understanding when different distributional representations succeed or fail—is arguably more useful to practitioners than the minimax optimality result, which guarantees good worst-case behavior but says nothing about instance-dependent performance differences between algorithms. It opens a direction for adaptive or hybrid methods that select representations based on environment characteristics inferred from initial samples.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The experiments use four manually constructed Markov reward processes (MRPs), not a standard benchmark dataset. These are: (1) Chain: a 10-state chain x₁ ↔ x₂ ↔ … ↔ x₁₀ with equal transition probability to neighbors, terminal states at ends, reward of 1 only at state 10. (2) Low random: 5 states with transition probabilities drawn from Dirichlet(0.01, …, 0.01) (concentrating mass on few transitions, producing near-deterministic dynamics), rewards i.i.d. Unif([0, 1]) per state. (3) High random: same structure as Low random but with Dirichlet(10, …, 10) (producing diffuse, high-entropy transitions). (4) Two-state: a 2-state MRP with transition matrix [[0.6, 0.4], [0.8, 0.2]] from Rowland et al. [2024, Example 6.5], where state 1 has reward 0 and state 2 has reward 1. These environments are chosen to span a range of structural properties: chain-like sparse transitions, near-deterministic vs. highly stochastic dynamics, and environments with tight self-loops known to challenge distributional RL methods. All experiments are fully tabular—there is no function approximation and no train/test split; the MRP parameters are synthetically generated once (for the random environments, a single random seed fixes the transition and reward matrices). Return distributions are computed for evaluation in all states, and ground-truth return distributions are estimated via first-visit Monte Carlo with at least 10⁴ returns per state and a truncation horizon T ensuring γᵀ/(1−γ) < 10⁻⁴. The paper reports no benchmark comparisons on standard RL evaluation suites (e.g., Atari, MuJoCo, DeepMind Control Suite), consistent with its tabular, model-based focus.
-
Base model(s). The paper does not use a neural network base model or pretrained LLM—all experiments operate on the four tabular MRPs described above. The "model" is the environment's true transition kernel P, which the algorithm estimates from samples. The generative model provides N i.i.d. next-state samples from each state (N ∈ {10², …, 10⁶}), and the algorithm constructs empirical transition probabilities P̂ via frequency counting. This is a tabular, non-parametric setting: the only "model capacity" parameter is m (the number of atoms/categories for distributional RL methods), not a neural architecture choice.
-
Metrics. The primary metric is supremum-Wasserstein-1 distance between the estimated return distributions and the Monte Carlo ground truth: max_{x∈𝒳} w₁(η̂(x), η^{true}(x)). This matches the theoretical analysis in Theorem 5.1, which bounds this exact quantity. A secondary metric is wallclock time (in seconds), used to compare computational efficiency across methods. For the sample-size scaling experiments (Figure 10), the paper also reports Wasserstein-1 error as a function of N, with 95% bootstrapped confidence intervals across 30 independent runs (each with different sampled transitions). The Wasserstein-1 distance is computed as the L₁ difference between CDFs integrated over the return interval [0, (1−γ)⁻¹]—for categorical methods, this is computed exactly from the estimated CDF vectors; for QDP, it is computed from the quantile representation. There is no separate validation metric or held-out evaluation—all reported error is against ground truth for the same MRP instances.
-
Baselines. The paper compares four methods:
- DCFP (the proposed algorithm): solves the linear system (I − T̃_P̂)F̃ = H̃ directly using a sparse linear solver to obtain the exact categorical fixed point. Implemented in two variants: a sparse implementation exploiting the structure of T_P̂ (denoted DCFP), and a dense implementation using standard NumPy linear algebra (denoted d-DCFP) for reference.
- CDP (categorical dynamic programming): iteratively applies η ← Πₘ𝒯η with the empirical operator, as described by Bellemare et al. [2017] and Rowland et al. [2018]. Run for sufficiently many iterations (30,000) to ensure approximate convergence to the same fixed point as DCFP. Also implemented in sparse (CDP) and dense (d-CDP) variants.
- QDP (quantile dynamic programming): as introduced by Dabney et al. [2018b] and analyzed by Rowland et al. [2024], represents return distributions as mixtures of m quantiles at adaptively-chosen locations, applying a quantile projection after each Bellman update. This is the primary alternative distributional RL method based on a fundamentally different representation (quantiles rather than categorical histograms).
- Majority voting is not applicable in this tabular distribution estimation setting, and no mean-return-only baseline (e.g., certainty equivalence from Azar et al. [2013]) is reported.
-
Generation budget / compute accounting. The "generation budget" is the number of samples per state N, varied across {10², 10³, 10⁴, 10⁵, 10⁶}. This is the standard accounting in the generative model literature: each state independently receives N next-state samples. For wallclock-time comparisons, all methods use the same N = 10⁶ samples. The categorical methods additionally have a computational budget parameter m (number of atoms), swept across {30, 100, 300, 1000}, and CDP has an iteration budget (30,000 iterations; the paper verifies this is sufficient for convergence). QDP uses the same m to denote number of quantiles (making the comparison fair per-parameter). There is no separate accounting for the cost of building the empirical transition matrix P̂ or constructing the T_P̂ operator—these are one-time costs absorbed into the total runtime.
-
Cross-validation / statistical protocol. There is no cross-validation, as this is a fully generative setting: the MRP parameters are fixed, the only randomness comes from the N transition samples, and evaluation is against the ground-truth return distributions computed via Monte Carlo. The paper runs 30 independent repetitions with different sampled transitions (different random seeds for the generative model) and reports means with 95% bootstrapped confidence intervals. For the wallclock-time experiments (Figures 6–9), all 30 repetitions are aggregated into the reported curves; the paper does not report per-seed variance on wallclock time, only on Wasserstein error. The discount factor γ is varied across {0.8, 0.9, 0.95, 0.99} to assess sensitivity to the effective horizon. For the categorical methods, the paper also compares two atom-location strategies: global return range (atoms spanning [0, (1−γ)⁻¹], the theoretical default) and environment-specific return range (atoms spanning [r_min(1−γ)⁻¹, r_max(1−γ)⁻¹] when the reward range is known a priori). The environment-specific variant is not part of the theoretical guarantees but is included to show practical improvements from tighter prior knowledge.
Main Quantitative Results
Wallclock Time vs. Accuracy Trade-offs (Figures 6–9)
The central empirical message is that no single method dominates across all environments, discount factors, and atom counts—the relative ranking depends on environment structure. The paper reports results for each of the four environments at four discount factors and four atom counts, with N = 10⁶ fixed, plotting Wasserstein-1 error against wallclock time.
Chain environment (Figure 6). With the global return range, QDP consistently achieves the lowest asymptotic Wasserstein-1 error across all (γ, m) configurations—for example, at γ = 0.9, m = 100, QDP reaches ~0.02 Wasserstein error while DCFP/CDP reach ~0.05–0.08. However, DCFP and CDP are substantially faster in wallclock time: at m = 1000, DCFP solves in ~10² seconds while QDP requires ~10³ seconds, a 10× speed advantage. The sparse implementations (DCFP, CDP) dramatically outperform their dense counterparts (d-DCFP, d-CDP) at high atom counts—at m = 1000 and γ = 0.99, sparse DCFP takes ~10² seconds vs. ~10³ for dense, confirming that exploiting the sparsity structure (≈ 2⌈2/γ⌉ non-zeros per row) is critical for scaling m. With environment-specific atom locations (Figure 6b), DCFP and CDP close the accuracy gap to QDP substantially—at γ = 0.9, m = 100, DCFP now reaches ~0.015 vs. QDP's ~0.012, and at higher m, the gap narrows further. The improvement is because the chain's true return range is [0, 1] (only the terminal state gives reward), much narrower than [0, (1−γ)⁻¹], so environment-specific atoms effectively increase resolution.
Low random environment (Figure 7). Here, QDP dominates in accuracy across nearly all settings—at γ = 0.9, m = 100, QDP reaches < 0.05 Wasserstein error while DCFP/CDP reach ~0.2–0.5. This environment has near-deterministic transitions (Dirichlet(0.01, …, 0.01) concentrates probability mass on one or two next states), producing concentrated return distributions that QDP's adaptive quantile placement captures efficiently. The categorical methods' uniform grid wastes atoms on regions of the return interval with negligible probability mass. DCFP and CDP maintain their speed advantage (e.g., at m = 1000, ~10² vs. ~10³ seconds), but the accuracy gap is large enough that QDP may be preferred despite the runtime cost. Environment-specific atom ranges (Figure 7b) improve DCFP/CDP modestly but do not close the gap—the issue is not the range width but the distribution's concentration within that range.
High random environment (Figure 8). The pattern shifts: QDP and categorical methods are much closer in accuracy. At γ = 0.9, m = 100, QDP reaches ~0.05, DCFP reaches ~0.1. As m increases to 1000, the gap narrows further—DCFP reaches ~0.03 vs. QDP's ~0.02. The diffuse transitions in this environment (Dirichlet(10, …, 10) spreads probability mass broadly) produce return distributions that are themselves diffuse, reducing QDP's adaptive-placement advantage—both methods need to cover the return interval broadly, and the uniform grid is less wasteful. DCFP again holds a significant wallclock time advantage.
Two-state environment (Figure 9). This is the most striking reversal. At γ = 0.99 and m = 30, QDP with global return range achieves ~1.0 Wasserstein error (essentially useless), while DCFP achieves ~0.5. As m increases, QDP improves (reaching ~0.2 at m = 1000), but DCFP consistently outperforms it (reaching < 0.1 at m = 1000). This environment has tight self-loops—state 1 transitions to itself with probability 0.6, state 2 with 0.2, creating sharp bootstrap-induced features in the return distribution. Rowland et al. [2024] previously identified that QDP's fixed-point approximation can be poor in such environments because the quantile projection interacts poorly with the bootstrap structure. DCFP, by contrast, benefits from its uniform coverage—it doesn't try to concentrate atoms where mass is highest, so it avoids this systematic bias. The sparse implementations again provide ~10× speedup at high m.
DCFP vs. CDP. Across all environments, DCFP and CDP achieve nearly identical asymptotic Wasserstein errors (as expected, since they converge to the same fixed point given sufficient CDP iterations). The difference is wallclock time: at γ = 0.99, CDP requires ~30,000 iterations to converge, making it ~5–10× slower than DCFP's direct linear solve. This gap narrows at lower discounts (e.g., γ = 0.8, where √γ ≈ 0.89 converges faster). The direct-solve advantage is most pronounced at high γ, precisely where CDP's slow contraction (factor √γ ≈ 0.995, requiring ~4600 iterations for 10⁻⁴ error per the paper's calculation) becomes a practical bottleneck.
Sample Size Scaling (Figure 10)
The paper examines how Wasserstein-1 error decreases as N (samples per state) increases from 10² to 10⁶, for m ∈ {30, 100, 300, 1000} and γ ∈ {0.8, 0.9, 0.95, 0.99} in three environments (Chain, Low random, High random). The two-state environment is omitted from this analysis.
General scaling behavior. As expected, all methods improve with more samples—the Wasserstein error curves decrease roughly as a power law in N, consistent with the theoretical N⁻¹/² rate from Bernstein-style concentration. The improvement is most pronounced in the Chain environment (where the error drops from ~0.5–1.0 at N = 10² to ~0.01–0.05 at N = 10⁶ for DCFP with m = 1000) and least pronounced in Low random (where QDP already achieves low error at N = 10² due to its representation advantage).
Dependence on m. For categorical methods (DCFP shown, CDP identical), increasing m consistently reduces error at all N levels—no evidence of overfitting or statistical penalty from more parameters. This is consistent with Theorem 5.1's independence of N from m (beyond the m ≥ 4(1−γ)⁻²ε⁻² + 1 threshold). At N = 10⁶, the improvement from m = 30 to m = 1000 is dramatic: in Chain with γ = 0.9 and environment-specific atoms (Figure 10b), DCFP error drops from ~0.1 to ~0.01. For QDP, increasing m also reduces error, but with diminishing returns in some configurations—in High random at γ = 0.99 with 30 atoms, QDP performs anomalously (Figure 10a), showing error that increases slightly between N = 10² and N = 10⁶, suggesting that 30 quantiles are insufficient to capture the return distribution at high discount in this stochastic environment.
Global vs. environment-specific atom ranges (Figure 10a vs. 10b). In the Chain environment, environment-specific atoms provide a massive improvement—at γ = 0.9, m = 30, DCFP global-range error is ~0.3–0.5 across N, while environment-specific error is ~0.05–0.1, a ~5× reduction. This is because the global range [0, (1−γ)⁻¹] at γ = 0.9 spans [0, 10], but the true return range is [0, 1], so 29 of 30 atoms cover irrelevant values. In Low and High random, the improvement is smaller because the reward range [0, 1] is closer to the full return range, reducing the waste.
Anomalous QDP behavior. At N = 10² in High random with γ = 0.99, QDP with m = 30 achieves lower error than at N = 10⁶ (Figure 10a, green curves). The paper does not comment extensively on this, but it is consistent with QDP's non-linear fixed point being sensitive to the empirical transition matrix's structure at small sample sizes—a phenomenon that would not occur for the linear categorical methods.
Ablation Studies and Robustness Checks
The paper does not contain formal ablation studies in the conventional sense (varying one component of the algorithm while holding others fixed). However, several experimental variations serve the function of ablations by isolating the effects of specific design choices:
Sparse vs. dense implementation (Figures 6–9, d-DCFP/d-CDP vs. DCFP/CDP): The sparse implementations exploit the structure of T_P̂ (at most 2⌈2/γ⌉ non-zeros per row for the atom-index dimensions, plus sparsity from P̂ for the state-index dimensions) and use SciPy's sparse matrix-vector multiplication and linear solvers. The dense implementations use standard NumPy dense linear algebra. Across all environments and atom counts, the sparse variants achieve identical accuracy (same fixed point) with substantially lower wallclock time. At m = 1000, γ = 0.99 in the Chain environment (Figure 6a), sparse DCFP takes ~10² seconds while dense DCFP takes ~10³ seconds—a ~10× speedup. This confirms that the sparsity analysis in Appendix G.3 translates to practical computational gains. The speedup is particularly important for the high-m regime needed for tight accuracy guarantees at high discounts.
DCFP direct solve vs. CDP iterative convergence (Figures 6–9, DCFP vs. CDP): Both methods converge to the same fixed point (their Wasserstein errors are identical to within experimental precision across all configurations), validating that 30,000 CDP iterations suffice for convergence. The wallclock-time gap—DCFP typically 5–10× faster than CDP at γ = 0.99—quantifies the benefit of avoiding slow DP iterations. At γ = 0.8, the gap is smaller (CDP converges faster due to √γ ≈ 0.89), consistent with CDP requiring O((1−γ)⁻¹ log(1/ε)) iterations while DCFP's linear solve cost is relatively insensitive to γ (it depends primarily on the sparsity pattern, not the contraction rate). This ablation justifies the paper's claim that DCFP is practically preferable to CDP for high-discount settings.
Global vs. environment-specific atom locations (Figures 6–8, panels (a) vs. (b); Figure 10): This is the closest the paper comes to an ablation of the categorical grid definition. In the Chain environment, environment-specific atoms (spanning [0, 1] instead of [0, (1−γ)⁻¹]) provide dramatic error reductions—at γ = 0.9, m = 100, DCFP error drops from ~0.05 to ~0.015 (Figure 6). In Low random, the improvement is modest (Figure 7), and in High random, it is intermediate (Figure 8). This quantifies how much of the categorical methods' disadvantage vs. QDP is attributable to wasted atoms outside the true return range. The finding suggests that in practice, any prior knowledge about the return range should be used to tighten the atom support—the theoretical analysis's use of the full [0, (1−γ)⁻¹] interval is a worst-case bound, not a practical recommendation.
Varying m and N jointly (Figure 10): The paper sweeps m ∈ {30, 100, 300, 1000} at each N ∈ {10², …, 10⁶}, providing an implicit ablation of the representation-vs-statistical-error tradeoff. The key observation: increasing m never increases error at any N, even at N = 10² where statistical error dominates. If there were a bias-variance tradeoff (more parameters → higher variance), we would expect optimal m to grow with N. The absence of such a tradeoff is consistent with the theoretical finding that N does not depend on m—the Cramér distance integrates over atoms, smoothing out per-atom variance. This is a non-trivial empirical validation of the paper's central statistical claim.
Varying γ (Figures 6–10): The four discount factors {0.8, 0.9, 0.95, 0.99} span a range where the effective horizon (1−γ)⁻¹ varies from 5 to 100. The results show that DCFP's wallclock-time advantage over CDP grows with γ (as expected from contraction theory), while QDP's accuracy advantage over DCFP in Low random is largely independent of γ. In the Two-state environment, QDP's disadvantage vs. DCFP is most pronounced at γ = 0.99 (Figure 9, m = 30: QDP error ~10 vs. DCFP ~0.5), confirming that the quantile fixed-point approximation error worsens with tighter bootstrapping.
Environments as ablations of structural properties: The four environments are not a random sample but are chosen to ablate specific structural features:
- Chain (sparse transitions, narrow return range): tests how methods handle limited reward propagation.
- Low random (near-deterministic, concentrated distributions): tests whether methods can exploit distributional concentration.
- High random (diffuse transitions, broad distributions): tests robustness to high stochasticity.
- Two-state (tight self-loops): tests behavior under the bootstrapping structure known to challenge QDP [Rowland et al., 2024].
The pattern of results validates that these structural features indeed matter: QDP excels in Low random (concentrated distributions), DCFP excels in Two-state and is competitive in High random (diffuse distributions + bootstrapping), and both methods perform reasonably on Chain with environment-specific atoms. No single method dominates across all these structural ablations, supporting the paper's implicit message that distribution representation choice is environment-dependent.
Critical Assessment
Does the empirical evaluation support the paper's central theoretical claim?
The paper's flagship theoretical result (Theorem 5.1) states that DCFP with m ≥ 4(1−γ)⁻²ε⁻² + 1 and N = Ω̃(ε⁻²(1−γ)⁻³) achieves ε-accurate Wasserstein-1 estimation with high probability. The empirical evaluation provides qualitative support for the directions predicted by the theory—error decreases with N and m—but does not constitute a rigorous empirical validation of the sample complexity bound. Specifically:
-
The experiments do not test the predicted scaling exponents. To validate (1−γ)⁻³ scaling, one would need to hold N and ε fixed while sweeping γ, fitting the exponent of (1−γ)⁻¹ in the achieved error. The experiments instead sweep N, m, and γ independently but do not attempt to extract scaling laws. The paper does not plot ε vs. N on a log-log scale or fit power laws, which would be the standard approach to validating a sample complexity rate.
-
The environments are too small to test asymptotic regimes. With |𝒳| ∈ {2, 5, 10}, the logarithmic dependence on |𝒳| (polylog(|𝒳|/δ)) is negligible—these experiments cannot distinguish polylog(|𝒳|) from log(|𝒳|) or from constant. The sample complexity bound's dependence on |𝒳| is essentially untested.
-
N ranges from 10² to 10⁶, but ε is never explicitly computed against the theoretical bound. The paper does not check whether, for a given N, γ, and m, the achieved ε matches (within constants) what Theorem 5.1 predicts. This would require instantiating the concentration constants (C_B1, C_B2, etc.) with their explicit logarithmic dependencies, which the paper does not attempt. The experiments are better described as validating the monotonicity of the error with respect to N, m, and γ—error goes down with more samples and more atoms, and problems get harder as γ → 1—rather than validating the precise rate.
-
The finite-m condition is not stress-tested. Theorem 5.1 requires m ≥ 4(1−γ)⁻²ε⁻² + 1. At γ = 0.99 and ε = 0.01, this demands m ≥ 4 × 10⁸—far beyond the m = 1000 maximum tested. The experiments operate in a regime where representation error likely dominates (m is too small to achieve the ε that the sample size N permits), making it a test of the categorical representation's approximation quality at small m, not of the full "choose m large enough, then the sample complexity kicks in" logic. The paper's empirical finding that "increasing m always helps" is consistent with operating below the m-threshold where representation error becomes negligible—in the theoretical regime, once m exceeds the threshold, further increases should yield diminishing returns (because the statistical error floor, not the representation error, dominates). The experiments never reach this saturation point.
What do the experiments actually demonstrate?
The experiments are more accurately described as a comparative study of distributional RL algorithms in tabular settings, evaluating practical trade-offs between representation choices (categorical vs. quantile) and computational strategies (direct solve vs. DP iterations). Their strengths are:
-
Clear documentation of the wallclock-time advantage of DCFP over CDP. This is a genuine practical finding: at high discounts, solving the linear system is much faster than iterating DP to convergence. The sparse implementation amplifies this advantage.
-
Identification of environment-dependent winners between categorical and quantile methods. The experiments reveal that QDP's adaptive support placement helps in near-deterministic, concentrated-distribution settings (Low random) but can hurt in environments with sharp bootstrap-induced features (Two-state). This is a practical diagnostic not captured by worst-case minimax theory.
-
Demonstration that environment-specific atom locations substantially improve categorical methods. The gap between global and environment-specific variants in the Chain environment (Figure 6) quantifies the benefit of incorporating prior knowledge about the return range.
Weaknesses and missing experiments
No comparison to mean-return-only baselines. The paper's central theoretical claim is that distributional RL achieves the same sample complexity as mean-return estimation. The experiments never compare distributional methods against a mean-return baseline (e.g., certainty equivalence from Azar et al. [2013]) to see whether, empirically, the distributional methods actually achieve comparable accuracy-per-sample. This would require defining a fair metric—perhaps the Wasserstein-1 error of the distributional method vs. the Wasserstein-1 error of a Dirac-delta-at-mean estimate—and is not attempted.
No test on standard benchmark MDPs. The four synthetic MRPs with |𝒳| ≤ 10 are far from realistic scale. The computational claims about sparse linear solvers would be more convincing on larger environments (|𝒳| ~ 10²–10³) where the contrast with dense methods is starker and where CDP's iteration count becomes more burdensome. Environments from the tabular RL literature (e.g., Garnet problems, random MDPs at scale, or classic domains like Frozen Lake with larger grids) would strengthen the empirical case.
No comparison to Zhang et al. [2023]. Given that the paper is framed as resolving an open question from Zhang et al. [2023], an empirical comparison to their (approximate) unprojected distributional DP would be natural. The paper acknowledges that Zhang et al.'s method is not practical without approximations, but implementing those approximations and comparing against DCFP would ground the theoretical improvement in empirical terms.
No statistical uncertainty quantification on the DCFP estimates. DCFP produces point estimates of the categorical fixed point; it does not provide confidence bands or credible intervals for the return distributions. The SC-CDF Bellman equation (Section 5.2) is introduced as encoding statistical fluctuations, but this is used only theoretically (to bound Σ_P in the proof), not to compute empirical uncertainty estimates. A natural extension would be to sample from the SC-CDF fixed point (by repeatedly applying single-sample operators) to produce bootstrap-like confidence bands, and to report whether these bands are well-calibrated. The experiments do not pursue this.
QDP implementation details matter but are not analyzed. QDP's performance depends on implementation choices: how quantiles are initialized, how the quantile projection handles edge cases, convergence criteria. The paper does not discuss these details or ablate them, treating QDP as a fixed baseline. Given that some of the results show QDP underperforming (Two-state) or behaving non-monotonically in N (High random, γ = 0.99, m = 30 in Figure 10), understanding whether these are fundamental limitations of quantile representations or artifacts of the specific QDP implementation would strengthen the comparative analysis.
The generative model assumption is not challenged. All experiments draw exactly N i.i.d. samples per state—the ideal generative model setting. The paper does not test robustness to violations: what if sample sizes vary across states (as in offline data collected by a behavior policy)? What if samples within a state are correlated (as in trajectory data)? These are the practical settings where model-based distributional RL would be deployed, and the theoretical guarantees do not automatically extend.
Do the experiments support the paper's positioning that distributional RL is "no harder" than mean estimation?
The experiments do not address this claim directly. They show that DCFP can achieve good accuracy, but they never benchmark the sample size required for a target Wasserstein-1 error against the sample size required for a target L∞ error in mean estimation on the same environments. This comparison would require running Azar et al.'s [2013] certainty equivalence algorithm on the same MRPs, measuring its value-function error as a function of N, and comparing to DCFP's distributional error—converting between metrics in a principled way. The paper does not attempt this, so the empirical section is silent on the "no harder than mean estimation" claim, which remains a purely theoretical result.
6. Limitations and Trade-offs
The Difficulty Estimation Cost Is Not Accounted For
The assumption or constraint. The compute-optimal framework requires estimating each prompt's difficulty before deciding how to allocate the inference budget. The paper's method for doing so—generating 2048 complete solutions per question and averaging either ground-truth correctness (oracle bins) or the PRM's final-answer score (predicted bins)—is extraordinarily expensive. As the authors acknowledge in Section 3.2:
"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"
This means the difficulty estimation step alone consumes 2048 generations per question—more compute than the largest test-time budgets studied (256–512 generations). The reported 4× efficiency gains over best-of-N are computed after difficulty is known, without amortizing the cost of learning it.
The consequence. In a realistic deployment where difficulty is unknown a priori, the total cost would be: (2048 generations for difficulty estimation) + (the compute-optimal strategy's budget). This could easily make the "compute-optimal" approach more expensive than simply running best-of-N with a large fixed budget on every question. For example, if the typical question is easy and the compute-optimal policy selects 4 generations of sequential revisions, the total cost is 2048 + 4 = 2052 generations—far worse than best-of-256 (256 generations). The 4× efficiency claim is therefore an upper bound on achievable efficiency that would only be realized if difficulty estimation can be made much cheaper, not a realized deployment gain.
What evidence exists in the paper. The paper acknowledges this explicitly (Section 3.2) but does not quantify the cost or include it in any budget calculation. Figures 4 and 8, which show the 4× improvements, all assume difficulty is known at zero cost. The paper does not compare total cost (estimation + execution) against baselines, nor does it report what fraction of the total compute budget difficulty estimation would consume under realistic assumptions.
Mitigation status. The paper suggests future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8) or developing adaptive schemes that estimate difficulty during the problem-solving process itself. However, no such method is developed or evaluated. This remains the single largest gap between the paper's theoretical framework and practical deployability.
The DCFP Algorithm Requires Solving a Large Linear System and Has No Theoretical Guarantees for Large State Spaces
The assumption or constraint. DCFP computes the categorical fixed point by solving the linear system (I − T̃_P̂)F̃ = H̃, where T̃_P̂ is an |𝒳|(m−1) × |𝒳|(m−1) matrix. The paper's theoretical sample complexity bounds (Theorem 5.1) hold for any finite state space |𝒳|, but the computational complexity of solving this linear system scales as O(|𝒳|³m³) if done densely. While the paper notes that sparsity can be exploited (Appendix G.3), there are no formal computational complexity guarantees—no bound on the time or memory required to solve the system as a function of |𝒳|, m, and the sparsity pattern. The empirical results (Section 6) use toy environments with |𝒳| ≤ 10 states, where the linear system has at most 10m unknowns. For a typical tabular RL problem with |𝒳| = 10⁴ states and m = 1000 (needed for ε = 0.01 at γ = 0.99 per Theorem 5.1), the linear system would have ~10⁷ unknowns—a sparse direct solve at this scale is far from trivial.
The consequence. The statistical optimality proven in Theorem 5.1 is decoupled from computational tractability. The theorem says N = Ω̃(ε⁻²(1−γ)⁻³) samples suffice statistically, but it does not bound the wallclock time or memory needed to actually compute the estimate from those samples. For practitioners, knowing that the sample complexity is optimal is cold comfort if the subsequent computation is infeasible. This limitation is particularly acute at high discounts (γ → 1), where m must grow as (1−γ)⁻² to control representation error, causing the linear system dimension to explode.
What evidence exists in the paper. The experiments (Figures 6–9) show wallclock times for |𝒳| ∈ {2, 5, 10} and m ≤ 1000—hardware-scale problems where dense solvers would already be slow (the paper uses sparse solvers to achieve ~10² seconds at m = 1000, |𝒳| = 10). But these are far from the regime where the theoretical guarantees are interesting (ε small, γ high, |𝒳| large). The paper does not report scaling behavior with |𝒳|—all experiments fix |𝒳| and vary m, N, γ. There is no experiment that tests whether DCFP's runtime remains practical as |𝒳| grows to hundreds or thousands of states.
Mitigation status. The paper acknowledges that dense solving costs O(|𝒳|³m³) (Section 4.1) and notes that sparsity can help, but provides no formal complexity bounds, no discussion of iterative linear solvers as alternatives to direct solves, and no experiments at scale. The purely theoretical discussion of sparsity in Appendix G.3 (bounding the number of non-zeros per row by 2⌈2/γ⌉) addresses only the m-axis sparsity, not the state-axis sparsity from P, and does not translate to a solvability guarantee. This limitation is largely unaddressed—the "direct" in DCFP is a conceptual advantage over iterative CDP, but both are computationally limited by the size of the linear system.
The Empirical Evaluation Only Tests on Four Small Synthetic Environments with No Comparison to Mean-Return Baselines
The assumption or constraint. All experiments in Section 6 and Appendix G use four synthetic Markov reward processes with state-space sizes |𝒳| ∈ {2, 5, 10}, random or structured transition matrices, and known reward functions. There are no experiments on standard RL benchmarks, no environments with function approximation, and—critically given the paper's central theoretical claim—no comparison to mean-return-only baselines such as the certainty equivalence algorithm of Azar et al. [2013]. The paper's headline message is that "distributional RL is essentially no harder, statistically speaking, than learning a value function" (Section 7), but the experiments never run a mean-return estimator alongside DCFP to check whether the sample efficiency is genuinely comparable in practice.
The consequence. The experimental section does not validate the paper's core theoretical claim empirically. To test whether distributional RL achieves the same accuracy-per-sample as mean estimation, one would need to: (1) run a mean-return baseline (e.g., certainty equivalence) on the same environments, (2) measure its L∞ value-function error as a function of N, (3) measure DCFP's Wasserstein-1 error as a function of N, and (4) compare the sample sizes needed to achieve equivalent quality (converting between metrics appropriately). The paper does none of this. The experiments instead compare DCFP against other distributional RL algorithms (CDP, QDP)—a useful relative comparison, but one that cannot speak to whether distributional RL's statistical cost matches mean-return RL's. Additionally, the four synthetic environments with |𝒳| ≤ 10 are far from the scale or diversity needed to assess whether the paper's structural findings (e.g., QDP winning on low-stochasticity, DCFP winning on tight self-loops) generalize to realistic problems.
What evidence exists in the paper. The experiments (Figures 6–10) demonstrate that DCFP's Wasserstein-1 error decreases with N and m, and that it outperforms QDP in specific environments (Two-state, high γ), but these are relative comparisons within the set of distributional methods. The paper reports no absolute sample-efficiency numbers that could be compared to published mean-return results, no scaling-law fits that extract exponents for (1−γ)⁻¹ or ε⁻², and no baseline that attempts to answer the question "given N samples, does DCFP estimate the return distribution as accurately as certainty equivalence estimates the mean?"
Mitigation status. The paper's empirics are explicitly positioned as "complement[ing] our theoretical analysis" (Section 6) and the authors do not claim the experiments validate the sample complexity rate. The theoretical contribution stands on its own proof. However, for practitioners who weigh empirical evidence heavily, the lack of a mean-return baseline and the small scale of the test environments significantly weaken the practical impact of the "no harder than mean estimation" message.
DCFP Works Only Under the Generative Model Access Regime, Which Is Rare in Practice and Bypasses the Exploration Problem
The assumption or constraint. The entire theoretical analysis—Theorem 5.1, the matching lower bound, the proof technique—assumes a generative model [Kearns et al., 2002, Kakade, 2003]: the algorithm can independently sample next-state transitions from any state on demand, receiving N i.i.d. draws per state. This is a clean theoretical setting because it isolates statistical difficulty from exploration difficulty, but it does not correspond to how most RL systems actually collect data. Practical data collection occurs through one of two mechanisms: online RL (the agent navigates between states sequentially, with sample sizes per state determined by the visitation distribution of its policy) or offline RL (the agent receives a fixed dataset of trajectories from an unknown behavior policy, with potentially highly imbalanced state coverage and correlated samples). Neither mechanism provides the uniform, independent per-state samples that DCFP's guarantees require.
The consequence. In an online setting, to collect N samples from each state, the agent must actually reach each state N times—which requires solving the exploration problem, potentially at an additional sample cost that scales polynomially in |𝒳| and (1−γ)⁻¹. The sample complexity would then multiply the statistical (1−γ)⁻³ factor by an exploration-dependent factor, likely at least (1−γ)⁻¹ or worse. The minimax optimality proven in Theorem 5.1 would not carry over; distributional RL might be strictly harder than mean estimation once exploration is accounted for. In an offline setting with non-uniform state coverage, the error guarantees would degrade in states with fewer samples in the dataset—the uniform max_{x∈𝒳} w₁ bound would become vacuous for poorly-covered states. The paper provides no guidance on how to adapt DCFP, or how its guarantees would change, under these more realistic data collection regimes.
What evidence exists in the paper. The paper is transparent about the generative model assumption throughout—it is stated in the abstract, Section 2.1, and Section 3. The experiments all faithfully implement the generative model regime (exactly N i.i.d. samples per state, controlled independently). There is no experiment that tests robustness to violations: no simulation with trajectory-based data collection, no offline dataset with imbalanced state coverage, no online learning scenario.
Mitigation status. The paper briefly discusses the policy optimization extension (Appendix F.2) under the assumption that the optimal policy is identified first (requiring additional generative model samples), and then its return distribution is estimated. This is within the same generative model framework. The paper does not address the substantially harder online or offline settings, does not suggest algorithmic modifications for these settings, and does not flag this as a central limitation—it is presented as a scope choice. For practitioners, this limitation means that DCFP as described is applicable primarily in settings with a simulator that allows arbitrary state reset (the closest real-world analog to a generative model), not in standard sequential RL tasks.
The Sample Complexity Bound Requires m ≥ 4(1−γ)⁻²ε⁻² + 1 Categories, Making the Representation Impractically Large at High Discounts
The assumption or constraint. Theorem 5.1 guarantees ε-accurate Wasserstein-1 estimation only when the number of categories m satisfies m ≥ 4(1−γ)⁻²ε⁻² + 1. This condition ensures that the representation error (the Cramér distance between the categorical fixed point F* and the true return distribution η*) is at most ε/2, so that the total error can be controlled by the statistical error alone. The required m grows as (1−γ)⁻², meaning that for problems with long effective horizons, the categorical grid must be extremely fine.
The consequence. For concrete parameter values: at γ = 0.99 and ε = 0.01, the condition demands m ≥ 4 × 10⁸ + 1 categories per state. This is manifestly infeasible—the linear system T̃_P̂ would have dimensions |𝒳| × (4 × 10⁸), and even storing the CDF vector for a single state would require ~400 MB of memory (at 8 bytes per float). For γ = 0.95 and ε = 0.05, the requirement is m ≥ 4 × 400 × 400 + 1 = 640,001—still extremely large for practical computation. This means the theoretical regime where the sample complexity bound applies (m chosen large enough to make representation error negligible) is computationally inaccessible for the very problems—high discount, high accuracy—where the bound is most interesting. Conversely, the experiments (Section 6) operate at m ≤ 1000, far below the theoretical threshold at γ = 0.99 and any reasonable ε—they are measuring performance in a regime where representation error likely dominates, not the regime the theorem describes.
What evidence exists in the paper. The experiments (Figures 6–10) sweep m ∈ {30, 100, 300, 1000}, never exceeding 1000. The observed error at γ = 0.99 with m = 1000 in the Two-state environment (Figure 9) is ~0.05–0.1 Wasserstein-1—far above the ε ~ 0.01 that would be needed to claim "high accuracy." The paper does not extrapolate what m would be needed to reach, say, ε = 0.01 at γ = 0.99 in these environments, nor does it compute the representation error bound (1−γ)⁻¹(m−1)⁻¹/² explicitly for its experimental configurations to verify that representation error is indeed dominated by statistical error. The theoretical condition m ≥ 4(1−γ)⁻²ε⁻² + 1 is stated but never empirically instantiated.
Mitigation status. The paper acknowledges this implicitly through its discussion of environment-specific atom locations (Appendix G), which can reduce the effective (1−γ)⁻¹ range and thus the required m for a given ε. For example, in the Chain environment where returns lie in [0, 1] rather than [0, (1−γ)⁻¹], the effective range is shorter and fewer atoms are needed. However, this mitigation only helps when the true return range is known a priori to be narrower than the worst-case [0, (1−γ)⁻¹]—for many environments, the full range is indeed possible, and the m requirement remains prohibitive. The paper does not address the fundamental tension: the bound that makes the sample complexity tight requires an m that makes the algorithm computationally infeasible. An algorithm that is both statistically optimal and computationally practical for high-discount problems remains an open challenge.
The Wasserstein-1 Metric May Not Capture What Practitioners Care About, and the Conversion from Cramér Distance Loses a (1−γ)^{-1/2} Factor
The assumption or constraint. The paper's theoretical guarantees bound the Wasserstein-1 distance between estimated and true return distributions: max_{x∈𝒳} w₁(η̂(x), η*(x)) ≤ ε. This metric is chosen because it is "particularly compatible" with categorical methods (Section 3) and upper-bounds differences in expectations of Lipschitz functions and risk measures like CVaR. However, the proof first establishes a bound in Cramér distance ℓ₂ (which has natural contractive properties for the categorical operator) and then converts via Lemma 5.2: w₁ ≤ (1−γ)^{-1/2} ℓ₂. This conversion introduces a (1−γ)^{-1/2} factor—the Cramér-distance bound of Ω̃(ε⁻²(1−γ)⁻²) becomes Ω̃(ε⁻²(1−γ)⁻³) in Wasserstein-1. A practitioner who cares about a different metric (e.g., Kolmogorov-Smirnov distance for tail bounds, or total variation for decision-making under uncertainty) would need a different conversion, potentially with an even worse dependence on (1−γ)⁻¹.
The consequence. The sample complexity's (1−γ)⁻³ scaling is partially an artifact of the chosen metric. If the goal were to bound Cramér distance directly, the sample complexity would be Ω̃(ε⁻²(1−γ)⁻²)—a full factor of (1−γ)⁻¹ better. This matters because Wasserstein-1 may not be the most practically relevant metric. For risk-sensitive applications, practitioners often care about tail behavior—the probability of exceptionally low returns, captured by value-at-risk (VaR) or conditional value-at-risk (CVaR). The paper notes that Wasserstein-1 upper-bounds CVaR differences, but this bound can be loose: two distributions can have small Wasserstein-1 distance while differing substantially in their 1% quantile. For example, if η* has a 1% chance of return 0 and η̂ has a 0% chance of return 0 but slightly more mass at low returns, the Wasserstein-1 distance could be small while the 1%-CVaR difference is large. The paper's guarantees do not directly address tail-risk estimation accuracy, which is often the primary motivation for using distributional RL over mean estimation.
What evidence exists in the paper. The paper uses Wasserstein-1 throughout the theoretical development (Theorem 5.1) and all experiments (Figures 6–10). There is no analysis of alternative metrics, no empirical evaluation of whether Wasserstein-1 error correlates with downstream decision quality, and no comparison of how sample complexity would change under different metric choices. Zhang et al. [2023] studied Kolmogorov-Smirnov and total variation metrics and found them to be statistically harder (requiring more samples)—the current paper does not engage with this finding or discuss whether the categorical approach offers advantages in these metrics.
Mitigation status. The paper's metric choice is explicitly motivated (Section 3) and the conversion from ℓ₂ to w₁ is mathematically clean (Lemma 5.2 is tight—Jensen's inequality gives equality when the CDF difference is constant). The authors do not claim their guarantees extend to stronger metrics. However, they also do not discuss the practical implications: a user deploying DCFP for risk-sensitive decisions would need to translate the Wasserstein-1 guarantee into a guarantee on the specific risk measure they care about, potentially introducing additional conservatism. This limitation is inherent to the metric choice and is not addressed, except through the brief note that Wasserstein-1 controls expectations of Lipschitz functions and CVaR.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper resolves a foundational open question in distributional reinforcement learning: Is estimating full return distributions statistically harder than estimating just their means? The answer—proven and not merely conjectured—is that in the generative model regime, the two tasks have the same minimax sample complexity, up to logarithmic factors. This is a closure result, not a paradigm shift, but its implications ripple through how the field should think about the statistical cost of distributional methods.
Before this work, the gap between Zhang et al.'s [2023] upper bound of Ω̃((1−γ)⁻⁴) and their lower bound of Ω̃((1−γ)⁻³) left open the possibility that distributional RL genuinely requires an extra factor of (1−γ)⁻¹ in sample size—a substantial penalty at high discount factors. Practitioners weighing whether to invest in distributional methods for risk-sensitive applications had no rigorous answer to "how many more samples will this cost me?" Theorem 5.1 closes this gap by exhibiting a practical, implementable algorithm (DCFP) that achieves the lower bound. The message is clear: there is no hidden statistical tax for estimating return distributions rather than mean returns, at least under the clean conditions of the generative model regime. This converts distributional RL from a method with uncertain statistical overhead into one whose sample complexity is precisely characterized and provably optimal.
The work also reframes how categorical distributional RL should be analyzed. The introduction of the stochastic categorical CDF Bellman equation (Section 5.2, Definition 5.6) provides the missing analytical bridge between the deterministic contraction theory of Rowland et al. [2018] and the stochastic concentration arguments needed for sample complexity bounds. This equation—which encodes the statistical fluctuations of repeated categorical DP updates under random sampling—is a reusable tool. Its structure (Proposition 5.11: Σ_Q ≥ σ_Q + γQ Σ_Q − correction) mirrors the variance Bellman equation that underpins mean-return sample complexity analysis [Azar et al., 2013], but with an additional projection-induced correction term that has no mean-return analog. Any future work that needs to propagate local estimation errors to global guarantees for categorical distributional RL—whether in online settings, with function approximation, or under different sampling assumptions—will need an analog of this equation. The paper has given the field a template for such analyses.
Equally important is what this work demystifies about the cost of rich representations. A natural worry when moving from scalar value functions to full distributions is that the increased representational capacity (m parameters per state rather than 1) would inflate sample complexity—either directly (more parameters to estimate) or through log(m) factors from union bounds over atoms. The proof shows this worry is misplaced: the Cramér distance integrates over the return interval, so per-atom estimation errors average out rather than accumulating. The sample complexity bound has no dependence on m (beyond the requirement that m is large enough to control representation error), and the proof avoids per-coordinate concentration inequalities in favor of a Hilbert-space Bernstein bound that treats the entire CDF as a point in a function space. This is a statistical diagnostic: the effective degrees of freedom in categorical distributional RL are controlled by the smoothness of the CDF (O(1)), not by the discretization granularity (m). The experiments (Figure 10) provide qualitative support: increasing m from 30 to 1000 never increases error at any fixed sample size N, consistent with the absence of a bias-variance tradeoff.
The work also provides a reconciliation of conflicting empirical signals about categorical versus quantile methods. Practitioners have reported mixed results when choosing between these representations—sometimes categorical wins, sometimes quantile. The paper's experiments (Section 6, Figures 6–9) provide a diagnostic framework: quantile methods (QDP) excel in near-deterministic environments with concentrated return distributions (Low random, Dirichlet(0.01) transitions), where their adaptive support placement captures relevant outcomes efficiently, while categorical methods (DCFP) are more robust in environments with diffuse distributions or tight bootstrapping loops (Two-state, high γ), where quantile fixed-point approximations can develop systematic biases. This is not a one-method-dominates story but an environment-dependent tradeoff, and the paper provides structural language for predicting which method will work where.
Finally, the paper redirects attention from algorithm design to representation choice. DCFP is not a novel algorithm in conception—it is the natural certainty-equivalent approach applied to categorical distributional RL. Its contribution is proving that this straightforward algorithm, when properly configured (with m chosen to control representation error), achieves optimal sample complexity. This suggests that the bottleneck for distributional RL is not algorithmic sophistication (e.g., better search strategies for the Bellman update) but rather the interface between representation and statistical analysis. The field's energy may be better spent on understanding how different representations (categorical, quantile, particle-based, moment-based) trade off approximation error against statistical efficiency, rather than on developing more complex DP procedures.
Follow-Up Research This Work Enables
Instance-dependent bounds for categorical vs. quantile distributional RL, explaining when each representation excels. The paper's experiments reveal striking environment dependence—QDP dominates in Low random (near-deterministic), DCFP dominates in Two-state with high γ (tight self-loops)—but the theoretical analysis provides only worst-case minimax guarantees that cannot distinguish these cases. A natural follow-up would derive instance-dependent sample complexity bounds for both categorical and quantile methods, expressed in terms of measurable environment properties (e.g., the smoothness of the return distribution, the spectral gap of the transition matrix, the concentration of the reward distribution). The SC-CDF Bellman equation from this paper provides the local-to-global propagation machinery for categorical methods; an analogous "stochastic quantile Bellman equation" would be needed for QDP. A strong result would be a theorem of the form: "If the return distributions have bounded density (Lipschitz condition), DCFP achieves error ε with N = O(f(γ, ε, |𝒳|, smoothness)); if additionally the return distributions are concentrated (small support relative to the full range), QDP achieves error ε with N = O(g(γ, ε, |𝒳|, concentration)) where g ≪ f." The experiments in Section 6 already provide the qualitative pattern; a theory that quantifies it would turn practitioner heuristics into rigorous guidance.
Computational complexity bounds for DCFP with sparse linear solvers on factored state spaces. The paper's sample complexity result (Theorem 5.1) is decoupled from computational cost—the linear system (I − T̃_P̂)F̃ = H̃ has dimension |𝒳|(m−1), and the paper provides only empirical runtime measurements on |𝒳| ≤ 10 states (Figures 6–9). A critical open question: for what classes of MDPs can this linear system be solved in time polynomial in |𝒳| and m? The sparsity analysis in Appendix G.3 bounds the number of non-zeros per row by 2⌈2/γ⌉ (from the atomic structure) plus the sparsity from P̂. For factored MDPs where the transition matrix has low rank or tensor structure, direct solvers may achieve near-linear time. A concrete research direction: implement DCFP on the standard "Garnet" random MDP benchmark at scale (|𝒳| = 10³–10⁴), measure wallclock time and memory as a function of |𝒳|, m, and γ, and determine the largest feasible problem size. If DCFP with sparse solvers scales to |𝒳| ~ 10⁴ at m ~ 100, it becomes practical for many tabular planning problems; if it breaks down at |𝒳| ~ 100, then iterative methods (CDP with early stopping, or stochastic approximations) become necessary despite their slower convergence. This experiment would ground the paper's computational claims in concrete scaling behavior.
Extending the minimax result to the online RL setting with exploration. Theorem 5.1 assumes a generative model—N i.i.d. samples per state, independent of the agent's behavior. Online RL, where the agent must navigate to states to collect samples, introduces two complications: (1) sample sizes per state are random and determined by the visitation distribution of the agent's policy, and (2) samples collected along a trajectory are correlated (the next state depends on the current state). A natural extension is to analyze DCFP-style distributional estimation in the online PAC RL setting: the agent interacts for T steps, builds an empirical model from the observed transitions (which have non-uniform state coverage), and returns estimated return distributions. The key question is whether the additional (1−γ)⁻¹ factor that typically appears in online mean-return PAC bounds (from the exploration cost) also appears for distributional estimation, or whether distributional methods suffer an additional penalty. The SC-CDF Bellman equation from this paper would be the starting point for the analysis, but it would need to be modified to handle the random, policy-dependent sample counts. A realistic goal: prove that with an appropriate exploration policy (e.g., UCB-style bonus), DCFP achieves ε-accurate distributional estimates in Wasserstein-1 distance with T = Ω̃(ε⁻²(1−γ)⁻⁴|𝒳|) steps, matching the online mean-return lower bound.
Bootstrapped confidence bands for return distributions via the SC-CDF Bellman equation. Section 5.2 introduces the SC-CDF Bellman equation as a theoretical tool for bounding the local-to-global propagation of estimation errors, but the paper never uses it for uncertainty quantification—producing confidence bands around the estimated return distributions. A natural follow-up: instead of solving the fixed-point equation only once to get point estimates F̂, repeatedly apply the single-sample operator T̂ (with fresh samples from the empirical model P̂) to generate samples from the SC-CDF fixed-point distribution Φ_P̂. These samples encode epistemic uncertainty: for each state x, they produce a distribution over possible CDF vectors, which can be summarized as pointwise confidence bands (e.g., the 5th and 95th percentiles of F_i(x) across samples). A concrete experiment: on the four environments from Section 6, generate these bootstrap-like confidence bands for DCFP estimates, check whether the true return distribution's CDF lies within the bands at the nominal coverage rate (e.g., 95%), and compare the band width to the theoretical error bound from Theorem 5.1. This would transform the SC-CDF Bellman equation from a proof device into a practical inference tool, giving practitioners not just point estimates of return distributions but also reliable measures of estimation uncertainty—a capability that no current distributional RL algorithm provides.
Distributional RL with function approximation: when does the categorical projection regularize? All of this paper's results are in the tabular setting. Extending to linear function approximation or neural networks raises a fundamental question: does the categorical projection Πₘ act as a regularizer that improves statistical efficiency, or does it introduce bias that hurts? In the tabular case, the paper shows that projection does not harm the sample complexity rate (it achieves the same (1−γ)⁻³ as unprojected methods [Zhang et al., 2023], but with a practical algorithm). In the function approximation case, the projection forces the predicted return distributions to lie in the span of the categorical basis, which could either reduce variance (by constraining the effective parameter count) or increase bias (by forcing a poor fit to the true return distributions). A concrete experiment: on a suite of continuous-state MDPs (e.g., classic control problems with discretized state or linear features), compare the sample efficiency of categorical distributional TD (with projection, à la Bellemare et al. [2017]) against unprojected distributional TD and mean-return TD, varying m and the feature dimension. The SC-CDF Bellman equation's structure suggests that the projection's smoothing effect (integrating over atoms) might provide implicit variance reduction that is particularly valuable in the low-data regime—a hypothesis that can be tested by measuring estimation error as a function of the number of trajectories.
Offline distributional RL with non-uniform coverage: guarantees for states with few samples. The generative model assumption ensures equal sample sizes across all states, which is unrealistic for offline RL where data comes from a behavior policy that visits some states frequently and others rarely. A practically important extension: given an offline dataset of K trajectories from a fixed behavior policy, what guarantee can DCFP provide for the return distribution at a state x that appears n(x) times in the dataset? The challenge is that the empirical transition matrix P̂ will have high variance for states with small n(x), and the error propagation through the Bellman equation couples all states—an error in a poorly-estimated state's distribution affects estimates at all states that can transition to it. A natural analysis would combine this paper's SC-CDF Bellman propagation with importance-weighting or pessimistic (penalized) estimation for low-coverage states, yielding a bound of the form: for a state x, the Wasserstein-1 error is bounded by ε with high probability provided the effective sample size (combining n(x) and the coverage of predecessor states) exceeds some threshold. The paper's experiments already provide a starting point—Figure 10 shows how DCFP's error scales with N (uniform across states). The natural follow-up would perturb this setup to give different states different N, mimicking offline data imbalance, and measure whether per-state error is predicted by per-state sample size or whether error "spills over" from poorly-sampled states.
Practical Applications and Downstream Use Cases
Risk-sensitive policy evaluation in simulated environments with state-reset capability. DCFP is directly applicable whenever a simulator provides generative-model-like access—the ability to independently sample transitions from arbitrary states. This setting arises in: (a) Medical treatment planning where a patient simulator (trained from electronic health records) can generate counterfactual outcomes for different treatment sequences. DCFP would estimate the full distribution of patient outcomes (e.g., time to recovery, probability of adverse events) under a candidate treatment policy, enabling risk-averse optimization (e.g., minimize the 10th percentile of recovery time rather than the mean). (b) Industrial process control where a digital twin of a factory floor can simulate equipment failure and throughput under different scheduling policies. DCFP provides the distribution of production metrics, allowing plant managers to guarantee throughput with high probability rather than just in expectation. (c) Algorithm configuration where a simulator evaluates candidate hyperparameter settings for an optimization algorithm across random problem instances. DCFP estimates the distribution of runtime or solution quality, enabling selection of configurations that are robust to worst-case problem instances. In each case, the paper's guarantee that N = Ω̃(ε⁻²(1−γ)⁻³) samples per state suffice for ε-accurate distribution estimation provides a principled budget: for a medical simulator with |𝒳| = 10³ states, γ = 0.99 (accounting for long time horizons), and target ε = 0.01 in Wasserstein-1, the theory predicts N ≈ 10⁸ samples per state are statistically sufficient (ignoring constants). Whether this is practically affordable depends on simulator speed, but the bound provides a starting point for cost-benefit analysis: if the simulator runs at 10⁴ samples/second, data collection would take ~10⁴ seconds (~3 hours) per state—prohibitive for 10³ states, but the bound is worst-case, and instance-dependent constants may be much smaller.
Model-based distributional RL as a post-hoc analysis tool after mean-return policy optimization. The paper's Appendix F.2 sketches how DCFP can be combined with policy optimization: first identify the optimal policy using mean-return RL (e.g., certainty equivalence with N₁ samples per state-action pair), then estimate its return distributions using DCFP (with N₂ additional samples). This two-phase approach is attractive when risk assessment is needed only after a policy is selected—the bulk of the sample budget is spent on efficient mean-return optimization (phase 1), and distributional estimation (phase 2) is performed only for the final policy. Theorem 5.1 guarantees that phase 2 requires N₂ = Ω̃(ε⁻²(1−γ)⁻³) samples to achieve ε-accurate distribution estimates. For a practitioner deploying an RL policy in a safety-critical setting (e.g., autonomous vehicle planning), this workflow provides: (a) the mean-optimal policy from phase 1, and (b) a verified distribution of outcomes (collision probability, comfort metrics, time-to-destination) with provable accuracy guarantees from phase 2. The key practical benefit over running distributional RL from the start is sample efficiency: the phase 1 optimizer can use the full action space (|𝒜| × |𝒳| samples) without the computational overhead of maintaining per-action return distributions, and phase 2 only evaluates the single chosen policy (|𝒳| samples). The paper's experiments (Figure 10) show that DCFP's error decays roughly as N^{-1/2} in practice, so halving ε requires 4× more samples—a predictable cost that can be budgeted.
Distributional model checking for sim-to-real transfer. When a policy trained in simulation is deployed in the real world, the sim-to-real gap must be quantified. DCFP provides a tool for this: (1) estimate return distributions under the simulator's dynamics (using the generative model with many samples), (2) estimate return distributions under real-world dynamics (using a limited number of expensive real-world rollouts as the "generative model" samples), and (3) compare the two distributions in Wasserstein-1 distance. Theorem 5.1 guarantees that with N real-world samples per state, the real-world distribution estimate is ε-accurate, so a detected Wasserstein-1 gap larger than 2ε can be confidently attributed to sim-to-real mismatch rather than estimation noise. The paper's experimental methodology—Monte Carlo ground truth for the "true" distributions, DCFP estimates from finite samples—is exactly this workflow in microcosm. For a robotics application where real-world rollouts cost 10⁸–$10⁹—clearly impractical. The practical value of DCFP in this setting would come from using it on a learned, low-dimensional latent state representation where |𝒳| is small, or from instance-dependent improvements that tighten the bound for typical environments.
When to Prefer This Method
The paper explicitly positions DCFP against both CDP (iterative categorical dynamic programming) and QDP (quantile dynamic programming), and the experimental and theoretical results provide clear tradeoff conditions:
-
Prefer DCFP over iterative CDP when: (a) the discount factor γ is high (γ ≥ 0.95), where CDP's √γ contraction rate requires thousands of iterations to converge, making DCFP's direct linear solve 5–10× faster in wallclock time (Figures 6–9, γ = 0.99); (b) an exact solution to the empirical fixed-point equation is required (e.g., for downstream theoretical guarantees), since CDP only converges asymptotically; (c) the state space is small enough (|𝒳| ≲ 10³) and m ≲ 10³ that the linear system fits in memory and a sparse direct solver completes in acceptable time. If the state space is large (|𝒳| ≳ 10⁴) and the linear system cannot be stored or solved directly, CDP remains preferable—each iteration is a sparse matrix-vector product, requiring only O(|𝒳|m × nnz_per_row) operations and O(|𝒳|m) memory for the current iterate.
-
Prefer DCFP over QDP when: (a) the environment has tight bootstrapping loops (states with high self-transition probabilities), as in the Two-state environment with γ = 0.99 (Figure 9), where QDP's quantile fixed-point approximation error is substantially larger than DCFP's categorical approximation error; (b) the return distributions are diffuse or multi-modal (High random environment, Figure 8), where QDP's adaptive quantile placement provides less advantage over DCFP's uniform grid, and DCFP's faster wallclock time (often ~10× at high m) becomes the deciding factor; (c) a guaranteed error bound in Wasserstein-1 distance is required, since Theorem 5.1 provides rigorous finite-sample guarantees while QDP's statistical properties are not yet characterized at this level of precision.
-
Prefer QDP over DCFP when: (a) the environment is near-deterministic with concentrated return distributions (Low random, Figure 7), where QDP's adaptive quantile support achieves substantially lower approximation error for a given atom count m; (b) the atom budget m is small relative to the return distribution's complexity, since QDP's support adaptation makes better use of each atom; (c) the goal is to minimize representation error (distance to the true return distribution) rather than to match a sample complexity bound, since DCFP's guarantees require m ≥ 4(1−γ)⁻²ε⁻² + 1—often impractically large at high γ.
-
Prefer environment-specific atom locations (for any categorical method) when: any prior knowledge about the return range is available. The Chain experiments (Figure 6) show that using the known range [0, 1] instead of the worst-case [0, (1−γ)⁻¹] reduces DCFP error by ~5× at a given m and N. In practice, if the reward function's bounds [r_min, r_max] are known, setting atom locations to span [r_min(1−γ)⁻¹, r_max(1−γ)⁻¹] is a zero-cost improvement that preserves the theoretical guarantees (since the true return distributions are still contained in the support).
These tradeoffs are grounded in the paper's specific experimental findings (Section 6, Figures 6–9) and theoretical analysis (Theorem 5.1 vs. the absence of comparable results for QDP). The paper does not frame this as a universal decision hierarchy but rather as environment-dependent guidance that practitioners should validate on their specific problem class.