Distributional Bellman Operators over Mean Embeddings

ArXiv: 2312.07358

🎯 Pitch

Return distributions can be updated in reinforcement learning using only simple linear algebra — no expensive imputation steps that decode sketches back into probability distributions. The sketch Bellman operator operates directly on mean embeddings, and the resulting Sketch-DQN agent runs faster than existing distributional deep RL baselines while matching their performance on Atari.

---

1. Executive Summary

This paper proposes a novel algorithmic framework for distributional reinforcement learning based on learning finite-dimensional mean embeddings of return distributions, introducing the sketch Bellman operator — a linear-algebraic update that propagates sketch values directly without requiring expensive imputation strategies that decode sketches back into probability distributions. The framework yields both dynamic programming (Sketch-DP) and temporal-difference learning (Sketch-TD) algorithms, evaluated on tabular Markov reward processes and scaled to deep RL on the Atari 57 benchmark, where a Sketch-DQN agent using sigmoidal mean embeddings approaches the performance of Implicit Quantile Networks (IQN) while running faster than QR-DQN and IQN. The paper provides asymptotic convergence theory establishing that Sketch-DP iterates converge to a neighborhood of the true mean embeddings, with the error bound shrinking inversely with the number of features — establishing that arbitrary accuracy is achievable as feature count increases, but only when the regression problem defining the Bellman coefficients can be solved with sufficiently low approximation error.

2. Context and Motivation

The Core Problem: Efficiently Representing Return Distributions Without Distributional I/O

Distributional reinforcement learning (RL) fundamentally asks a richer question than traditional value-based RL. Rather than learning only the expected return from each state — a single scalar value $V^\pi(x) = \mathbb{E}^\pi_x[\sum_{t=0}^\infty \gamma^t R_t]$ — distributional RL aims to learn the full probability distribution of the random return $\sum_{t=0}^\infty \gamma^t R_t$ for each state. This is valuable because the distribution encodes information that the expectation discards: variance, skewness, multi-modality, tail risk, and other properties relevant to risk-sensitive decision-making, exploration, and understanding the aleatoric uncertainty inherent in the environment.

However, representing and updating probability distributions algorithmically is fundamentally more difficult than propagating scalar values. Probability distributions are infinite-dimensional objects, and computational algorithms must operate with finite-dimensional representations. The central tension in distributional RL is how to choose a finite-dimensional representation that is simultaneously: (1) expressive enough to capture meaningful distributional structure, (2) closed under the distributional Bellman operator — or at least approximately closed — so that dynamic programming updates are well-defined, and (3) computationally efficient to update, particularly in the deep RL setting where millions of Bellman backups are performed during training.

The gap this paper addresses is a specific and important instance of this representation problem. Prior work had established two dominant paradigms for representing return distributions: direct parametric approximations of probability distributions (categorical, quantile-based), and an indirect approach based on learning statistical functionals of return distributions called sketches (Rowland et al., 2019). The sketch approach offers elegant theoretical properties — it allows one to specify which aspects of the return distribution to capture by choosing which functionals to learn — but it suffers from a critical practical bottleneck: each Bellman update requires an expensive imputation strategy that decodes sketch values back into approximate probability distributions, applies the standard distributional Bellman operator to those distributions, and then re-encodes the result back into sketch values. This imputation step introduces significant computational overhead and, as noted by neuroscientists, is biologically implausible as a model of distributional learning in the brain (Tano et al., 2020).

The paper's central conceptual move is to ask: Can we define a Bellman update that operates entirely in sketch space, bypassing the imputation bottleneck entirely? This would mean that the sketch values themselves — finite-dimensional vectors — are updated via a closed-form operator, without ever explicitly reconstructing a probability distribution. The answer, it turns out, is "approximately yes" for a much broader class of sketches than previously known, provided we are willing to accept a controlled approximation error that can be made arbitrarily small by increasing the sketch dimensionality.

Why This Gap Matters: The Practical and Theoretical Stakes

The imputation bottleneck in sketch-based distributional RL is not merely an inconvenience — it fundamentally limits the scalability and applicability of the approach along several dimensions.

Computational efficiency in deep RL. In deep RL settings like the Atari benchmark, agents perform millions of Bellman backups during training. Each backup in an imputation-based sketch method (SFDP/SFTD; Rowland et al., 2019) requires solving an optimization problem to reconstruct a distribution from sketch values. For a sketch with $m$ expectiles, this involves solving a convex quadratic program with constraints that the imputed distribution must be a valid probability distribution (simplex constraints). Even with efficient solvers, this per-update cost multiplies across millions of updates, making the approach significantly slower than direct distributional methods like C51 (Bellemare et al., 2017) or QR-DQN (Dabney et al., 2018b). The paper's experiments (Appendix D.4, Figure 11) quantify this: SFDP can be over 100× slower per dynamic programming iteration than the proposed Sketch-DP algorithm, which uses only matrix-vector products.

Biological plausibility. One of the exciting intersections of distributional RL is with neuroscience, where there is growing evidence that dopamine neurons encode something closer to a distribution over future rewards rather than a single expected value (Dabney et al., 2020; Lowet et al., 2020). The sketch framework — particularly the use of expectile-based sketches — has been proposed as a computational model for how neural populations might encode return distributions (Dabney et al., 2020). However, the imputation step is biologically implausible: neurons would need to explicitly reconstruct a probability distribution from their activity patterns, perform a full distributional Bellman backup, and re-encode the result — a level of explicit probabilistic computation for which there is no neuroscientific evidence. Tano et al. (2020) explicitly raised this concern, motivating the search for a more direct, local update rule that operates on the neural code itself. The sketch Bellman operator proposed here is a candidate for such a local rule: it takes the form $U(x) \leftarrow \mathbb{E}^\pi_x[B_R U(X')]$, a simple linear propagation that could plausibly be implemented by synaptic weight matrices.

Algorithmic generality. The SFDP framework requires designing a specialized imputation strategy for each choice of sketch — an expectile sketch needs an expectile-specific imputation, a quantile sketch needs a quantile-specific one, and so on. This fragments the algorithmic space and makes it difficult to experiment with novel sketches, since each new sketch requires solving a non-trivial inverse problem (reconstructing a distribution from the sketch values). The Bellman sketch framework, by contrast, provides a unified recipe: given any feature function $\phi : \mathbb{R} \to \mathbb{R}^m$, one precomputes Bellman coefficients $B_r$ by solving a least-squares regression problem (Equation 5), and then applies the same linear update rule regardless of the choice of $\phi$. This dramatically lowers the barrier to exploring new sketches, opening a rich design space of feature functions.

Theoretical understanding. Beyond practical concerns, the imputation-based approach creates a theoretical gap. The SFDP operator applies $\psi \circ \mathcal{T}^\pi \circ \iota$, where $\psi$ is the sketch, $\mathcal{T}^\pi$ is the distributional Bellman operator, and $\iota$ is the imputation strategy. Analyzing the composition of these three non-linear operations is mathematically cumbersome, and prior work did not provide error propagation bounds for SFDP. The Sketch-DP framework, by contrast, defines an operator $\mathcal{T}^\pi_\phi$ that acts linearly on sketch values, enabling the clean error propagation analysis presented in Section 4 — a novel convergence theory that bounds asymptotic error in terms of the sketch dimensionality and the regression quality of the Bellman coefficients.

Where Prior Approaches Fall Short

The paper identifies specific limitations in four classes of prior work, each of which motivates a different aspect of the proposed framework.

1. Direct distributional approximations (categorical and quantile). The categorical approach (Bellemare et al., 2017) represents return distributions as categorical distributions over a fixed grid of $m$ support points $\{z_1, \ldots, z_m\}$, with learnable probabilities $\{p_1, \ldots, p_m\}$. The quantile approach (Dabney et al., 2018b) instead fixes the probabilities (uniform weights $1/m$) and learns the particle locations $\{z_1, \ldots, z_m\}$. Both produce state-of-the-art deep RL agents, but they have fundamental limitations. Categorical methods require choosing the support grid in advance, and if the true returns fall outside this grid, the representation is necessarily distorted. Quantile methods optimize particle locations via the quantile regression loss, which involves non-linear operations (sorting, Huber loss computations). Critically, neither approach directly learns statistical functionals of the return — they learn to approximate the full distribution, which may be overkill if one only cares about certain properties (e.g., variance, tail expectiles). The sketch framework offers more targeted representation: choose which functionals matter and learn only those.

2. Imputation-based sketch methods (SFDP/SFTD). As extensively discussed above, the core limitation of Rowland et al. (2019) is the imputation step. The SFDP update takes the form $U \leftarrow \psi(\mathcal{T}^\pi(\iota(U)))$, where $\iota$ maps sketch values to distributions and $\psi$ maps distributions back to sketch values. For this to work, $\iota$ must approximate a pseudo-inverse of $\psi$: $\psi(\iota(u)) \approx u$. Designing such pseudo-inverses is non-trivial and sketch-specific. For expectile sketches, the imputation strategy requires solving a constrained optimization problem (Bellemare et al., 2023, Section 8.6). For quantile-based sketches, it requires careful interpolation. The computational cost and design complexity make it unattractive for large-scale applications and for rapid experimentation with new sketches. The paper's Figure 1 visually contrasts the SFDP and Sketch-DP updates, emphasizing how Sketch-DP eliminates the $\iota \to \mathcal{T}^\pi \to \psi$ round-trip through distribution space.

3. Moment-based approaches. Sobel (1982) studied exact dynamic programming for the first $m$ moments of the return, and Tamar et al. (2013; 2016) developed TD methods for the first two moments (mean and variance). These methods are tempting because moments are Bellman closed: the $m$-dimensional space of moment vectors is closed under the distributional Bellman operator, meaning exact DP is possible without approximation. However, the paper identifies a critical practical problem (Appendix D.1): as $m$ grows, the typical magnitudes of moments diverge dramatically. The $k$-th moment scales roughly as $(\text{return range})^k$, meaning that for $m=50$, the moments span dozens of orders of magnitude. This makes single-learning-rate TD methods catastrophically unstable — a learning rate appropriate for the first moment is far too large for the 50th moment, and vice versa. Figure 7 empirically confirms this: polynomial feature sketches with $m=50$ perform terribly across all tabular environments compared to sigmoid-based sketches of the same dimensionality. The paper's framework generalizes moment methods to a much broader class of feature functions that remain numerically well-behaved at high dimensions.

4. MMD-based methods (MMDRL). Nguyen-Tang et al. (2021) proposed learning particle locations $\{z_i\}_{i=1}^m$ by minimizing a maximum mean discrepancy (MMD; Gretton et al., 2012) between the current return distribution estimate and the bootstrapped target. This can be viewed as operating on mean embeddings in a reproducing kernel Hilbert space (RKHS) — the empirical distribution $\frac{1}{m}\sum_{i=1}^m \delta_{z_i}$ has a mean embedding in the RKHS given by $\frac{1}{m}\sum_{i=1}^m K(z_i, \cdot)$. However, the approach has significant limitations that the paper's framework resolves:

• No dynamic programming formulation. MMDRL's update is defined purely as a sample-based TD algorithm driven by MMD gradients. There is no corresponding DP operator, making it impossible to analyze in the clean operator-theoretic framework that the paper uses to prove convergence. The gradient-based update also does not naturally respect the linear structure of mean embeddings.

• Non-convex optimization over particle locations. The space of mean embeddings corresponding to uniform-weight particle mixtures $\{\frac{1}{m}\sum_{i=1}^m K(z_i, \cdot) : z_i \in \mathbb{R}\}$ is a non-convex subset of the RKHS. Optimizing over particle locations via gradient descent can get stuck in local minima, and there are no convergence guarantees for the TD algorithm. The authors note that "Nguyen-Tang et al. (2021) do not provide theoretical analysis for their algorithm" (Appendix B.7), and that the MMD contraction results they prove are about the true distributional Bellman operator, not about their approximate algorithm.

• Computational cost of kernel evaluations. Each MMDRL update requires evaluating the kernel $K$ between all pairs of current and target particles — an $O(m^2)$ operation that can be expensive for large $m$, and does not benefit from the precomputation of Bellman coefficients that Sketch-DP exploits.

The paper's framework sidesteps all three issues: it works in a finite-dimensional RKHS (spanned by the feature coordinates $\phi_1, \ldots, \phi_m$) where the mean embedding is simply a vector in $\mathbb{R}^m$; it defines a linear DP operator $\mathcal{T}^\pi_\phi$ that provides clean convergence theory (Section 4); and it precomputes Bellman coefficients $B_r \in \mathbb{R}^{m \times m}$ so that each update is a simple matrix-vector product costing $O(m^2)$ — with the $O(m^3)$ matrix inversion cost paid once upfront.

How This Paper Positions Itself: A Unifying Framework Through Linear Bellman Coefficients

The paper's core conceptual innovation is the Bellman coefficient — a matrix $B_r \in \mathbb{R}^{m \times m}$ defined for each possible immediate reward $r \in \mathcal{R}$ that approximately translates the feature function $\phi$ evaluated at a bootstrap return $r + \gamma g$ into a linear function of $\phi$ evaluated at $g$ alone:

$\phi(r + \gamma g) \approx B_r \phi(g) \quad \text{for all } g.$

When this relationship holds exactly, the sketch is Bellman closed, and the return distributions' mean embeddings satisfy their own exact linear Bellman equation: $U^\pi(x) = \mathbb{E}^\pi_x[B_R U^\pi(X')]$. The paper notes that Theorem 4.3 of Rowland et al. (2019) characterizes all Bellman-closed mean embedding sketches — they are limited to invertible linear combinations of the first $m$ moments, recovering the classical Sobel (1982) moment DP algorithms but inheriting their numerical instability at scale.

The key relaxation is to allow the relationship to hold only approximately, with $B_r$ defined as the solution to a least-squares regression problem (Equation 5) under a weighting distribution $\mu$ over returns. This expands the space of usable sketches enormously — including sigmoids, Gaussians, sinusoids, indicator functions, and any other set of basis functions — at the cost of introducing approximation error $\varepsilon_B$ in each Bellman backup. The convergence theory in Section 4 shows that this error can be bounded and, crucially, that the asymptotic error in the mean embedding estimates scales as $O(1/m)$ for appropriately designed sketches (Proposition 4.4), meaning the approximation becomes arbitrarily accurate as the sketch dimensionality increases.

This positions the paper at the intersection of several research threads while addressing their individual limitations:

• Relative to Rowland et al. (2019): It preserves the sketch-based philosophy (learn only the functionals you care about) while eliminating the imputation bottleneck and providing convergence guarantees.
• Relative to Sobel (1982): It generalizes exact moment DP to approximate DP with vastly better numerical properties at high dimensions.
• Relative to Nguyen-Tang et al. (2021): It provides a cleaner theoretical foundation with a well-defined DP operator, convergence analysis, and a computationally efficient precomputation-based implementation.
• Relative to direct distributional methods (C51, QR-DQN): It offers a new axis of algorithmic design — the choice of feature function $\phi$ — that is orthogonal to the categorical/quantile dichotomy and may capture distributional properties not easily expressed in either of those frameworks.

The paper also connects to the broader literature on mean embeddings and kernel methods. The sketch $\psi(\nu) = \mathbb{E}_{Z \sim \nu}[\phi(Z)]$ is precisely a mean embedding of the distribution $\nu$ into $\mathbb{R}^m$ (Smola et al., 2007; Sriperumbudur et al., 2010), with the corresponding reproducing kernel $K(z, z') = \langle \phi(z), \phi(z') \rangle$. This connection opens the door to importing techniques from the rich kernel methods literature — kernel selection, feature approximation, Nyström methods — into distributional RL, though the paper explores this only implicitly through its choice of feature families (the translation family in Equation 8).

3. Technical Approach

3.1 Reader Orientation

This is primarily a methodological and analytical paper that designs a new family of distributional reinforcement learning algorithms based on a simple but powerful insight: if we can approximately express a feature vector evaluated at a shifted-and-scaled return as a linear transformation of the feature vector evaluated at the original return, then the mean embeddings of return distributions satisfy their own approximate Bellman equation that can be solved with simple linear algebra. The system solves the problem of efficiently propagating sketch-based representations of return distributions through Bellman backups without the expensive imputation step that decodes sketch values into full probability distributions and back, by precomputing Bellman coefficient matrices $B_r$ for each possible reward $r$ via least-squares regression and then applying purely linear updates — matrix-vector multiplications — at each step of dynamic programming or temporal-difference learning.

3.2 Big-Picture Architecture (Diagram in Words)

The system has four major components connected in a pipeline:

1. Feature Map $\phi: \mathbb{R} \to \mathbb{R}^m$ — A user-chosen nonlinear function that maps scalar returns to $m$-dimensional vectors. This defines which aspects of the return distribution the sketch will capture. Common choices include sigmoids, Gaussians, sinusoids, or indicator functions arranged in a "translation family" over a grid of anchor points.

2. Bellman Coefficient Matrices $B_r \in \mathbb{R}^{m \times m}$ — Precomputed once before any RL algorithm runs, by solving a least-squares regression problem (Equation 5) that finds the matrix $B_r$ minimizing $\mathbb{E}_{G \sim \mu}[\|\phi(r + \gamma G) - B \phi(G)\|_2^2]$ under a weighting distribution $\mu$ over possible returns. This captures the approximate linear relationship $\phi(r + \gamma g) \approx B_r \phi(g)$ that makes sketch-space Bellman backups possible.

3. Sketch Bellman Operator $\mathcal{T}^\pi_\phi: (\mathbb{R}^m)^\mathcal{X} \to (\mathbb{R}^m)^\mathcal{X}$ — The core algorithmic engine. Given current sketch estimates $U: \mathcal{X} \to \mathbb{R}^m$, it produces updated estimates via $(\mathcal{T}^\pi_\phi U)(x) = \mathbb{E}^\pi_x[B_R U(X')]$ — a linear operation that propagates sketch values from successor states through the Bellman coefficients and averages over transitions. This replaces the SFDP round-trip of imputation $\to$ distributional Bellman backup $\to$ sketch extraction.

4. Algorithmic Instantiations — Two concrete algorithms built on $\mathcal{T}^\pi_\phi$: Sketch-DP applies the operator iteratively to all states (dynamic programming), and Sketch-TD applies it incrementally from sampled transitions using stochastic approximation: $U(x) \leftarrow (1 - \alpha) U(x) + \alpha B_r U(x')$.

Information flows as follows: feature map $\phi$ and weighting distribution $\mu$ are chosen $\to$ Bellman coefficients $B_r$ are precomputed via closed-form least-squares solution $B_r = C_r C^{-1}$ where $C = \mathbb{E}_{G \sim \mu}[\phi(G)\phi(G)^\top]$ and $C_r = \mathbb{E}_{G \sim \mu}[\phi(r + \gamma G)\phi(G)^\top]$ $\to$ sketch estimates $U$ are initialized (typically to $\phi(0)$ for all states) $\to$ Sketch-DP or Sketch-TD iteratively updates $U$ using only matrix-vector products with $B_r$ $\to$ optionally, imputation strategies reconstruct approximate distributions from the learned sketch values for visualization or downstream use.

3.3 Roadmap for the Deep Dive

• First, the mean embedding sketch and feature map design space — since the entire framework is parameterized by the choice of $\phi$, understanding this design space (Equation 2, the translation family in Equation 8) is prerequisite to everything else.
• Second, the Bellman coefficients — the core mathematical innovation. I'll derive the exact Bellman-closed case (Equation 3), explain why it's too restrictive (only moments), then show how the least-squares relaxation (Equation 5) generalizes to arbitrary $\phi$ while maintaining the linear structure, and derive the closed-form solution $B_r = C_r C^{-1}$.
• Third, the Sketch Bellman operator $\mathcal{T}^\pi_\phi$ and its properties — how the approximate linear relationship $\phi(r + \gamma g) \approx B_r \phi(g)$ enables the crucial step $\mathbb{E}^\pi_x[\phi(R + \gamma G(X'))] \approx \mathbb{E}^\pi_x[B_R \phi(G(X'))]$ and why linearity of $B_r$ is essential for exchanging expectation and transformation.
• Fourth, the Sketch-DP and Sketch-TD algorithms in full detail, including the precomputation phase, initialization, the iterative update rules, and the handling of practical considerations like unknown rewards and the constant feature dimension.
• Fifth, the convergence analysis framework — the error propagation structure (Proposition 4.2, Figure 3), the three error sources (Bellman approximation, reconstruction, embedding), and the concrete bound for indicator features (Proposition 4.4) showing $O(1/m)$ convergence.
• Sixth, the deep RL adaptation (Sketch-DQN) — how sketch values are predicted by neural networks, how value readout coefficients $\beta$ enable Q-learning, and the specific architectural choices that differ from QR-DQN.

3.4 Detailed, Sentence-Based Technical Breakdown

3.4.1 Mean Embedding Sketches and the Feature Map Design Space

What a mean embedding sketch is. A mean embedding sketch $\psi$ is defined by a feature function $\phi: \mathbb{R} \to \mathbb{R}^m$ and maps a probability distribution $\nu$ over returns to an $m$-dimensional vector:

$\psi(\nu) := \mathbb{E}_{Z \sim \nu}[\phi(Z)]$

where $\phi(Z) = (\phi_1(Z), \phi_2(Z), \ldots, \phi_m(Z))^\top \in \mathbb{R}^m$ and the expectation is taken coordinate-wise, so $\psi(\nu)_i = \mathbb{E}_{Z \sim \nu}[\phi_i(Z)]$ for each $i \in \{1, \ldots, m\}$.

What it computes: given a probability distribution $\nu$ over the real line, the sketch produces a finite-dimensional summary vector by evaluating $m$ different nonlinear functions $\phi_1, \ldots, \phi_m$ at random draws from $\nu$ and taking their expected values. Each coordinate of the sketch vector is one statistical functional of $\nu$ — the expectation of a particular basis function. The sketch $\psi(\nu)$ is a lossy summary: many different distributions can map to the same or similar sketch vectors, and the choice of $\phi$ determines which features of the distribution are preserved and which are discarded.

Why this form: the mean embedding is linear in the distribution: for any mixture distribution $\nu = \alpha \nu_1 + (1-\alpha) \nu_2$, we have $\psi(\nu) = \alpha \psi(\nu_1) + (1-\alpha) \psi(\nu_2)$. This linearity is what makes the sketch Bellman operator possible — it means that expectations over random transitions decompose nicely. If we instead used a nonlinear functional (e.g., the median or a quantile), we would not be able to propagate expectations through the Bellman equation in the same way.

The translation family of feature maps. The paper's primary design space for $\phi$ is the translation family (Equation 8):

$\phi_i(z) := \kappa(s(z - z_i)), \quad \forall i \in \{1, \ldots, m\}$

where $\kappa: \mathbb{R} \to \mathbb{R}$ is a base feature function (a nonlinear "activation" like sigmoid, Gaussian, or hyperbolic tangent), $s \in \mathbb{R}^+$ is a slope parameter controlling the sharpness of the features, and $\{z_1, \ldots, z_m\} \subseteq \mathbb{R}$ is a set of anchor points spread across the expected return range.

What it computes: each coordinate $\phi_i$ is a shifted and scaled copy of the same base function $\kappa$, centered at anchor $z_i$ and with "width" controlled by $1/s$. A large slope $s$ produces narrow, sharply-varying features that act like localized basis functions; a small slope produces wide, slowly-varying features that capture more global structure. The feature vector $\phi(z)$ at a return value $z$ represents how strongly $z$ "activates" each of the $m$ basis functions — for a Gaussian base feature with slope $s$, $\phi_i(z) = \exp(-s^2(z - z_i)^2/2)$ measures similarity between $z$ and anchor $z_i$ in a soft, kernel-like way.

Why this form: the translation family provides three tunable degrees of freedom — base function $\kappa$, number of features $m$, and slope $s$ — that control the richness-resolution tradeoff. Increasing $m$ adds more anchors, providing finer coverage of the return range and reducing approximation error (as the theory in Proposition 4.4 confirms). Adjusting $s$ controls the smoothness of the representation, which affects both the regression error in the Bellman coefficients and the ability of the sketch to capture sharp distributional features. The paper tests sigmoid ($\kappa(x) = 1/(1 + \exp(-x))$), Gaussian ($\kappa(x) = \exp(-x^2/2)$), parabolic ($\kappa(x) = 1 - x^2$ for $|x| \leq 1$, zero otherwise), hyperbolic tangent ($\kappa(x) = \tanh(x)$), sinusoidal, and indicator base features — a deliberately broad set to demonstrate the framework's generality.

Anchor placement and slope selection heuristics. The paper provides concrete heuristics (Appendix B.6) for choosing the anchor range and slope. The extreme anchor points are set to $\hat{G}_{\min} - a\hat{L}$ and $\hat{G}_{\max} + a\hat{L}$, where $\hat{L} = \hat{G}_{\max} - \hat{G}_{\min}$ is the estimated return range and $a \approx 0.4$ is a small buffer. The default slope is chosen so that features with "non-trivial support" (e.g., $[-2, 2]$ for sigmoid and Gaussian $\kappa$) overlap by 50% with adjacent features and cover the full return range — specifically, $s = 20/(\hat{G}_{\max} - \hat{G}_{\min})$ for sigmoid and Gaussian base features. This ensures that every point in the return range falls within the responsive region of multiple basis functions, preventing "dead zones" where the feature vector is nearly constant and thus uninformative about distributional structure.

The constant feature dimension. The paper notes (Remark 3.1, Appendix C.2) that the framework is essentially invariant to the choice of basis for the function space spanned by $\phi_1, \ldots, \phi_m$. In practice, one coordinate is often set to a constant (e.g., $\phi_1(g) \equiv 1$), which ensures that the sketch Bellman operator is affine (recovering the classical Bellman equation for expected values) rather than purely linear. The authors found this "very crucial for a good performance" in deep RL (Appendix C.2) and include it by default.

3.4.2 The Bellman Coefficients: Exact Closure, Relaxation, and Computation

The exact Bellman-closed case. The paper begins (Section 3) by considering the ideal scenario where the feature map satisfies an exact linear relationship under the bootstrap transformation:

$\phi(r + \gamma g) = B_r \phi(g)$

for all possible returns $g$ and each immediate reward $r \in \mathcal{R}$, with $B_r \in \mathbb{R}^{m \times m}$ not depending on $g$. This equation says: if we know the feature vector at a return value $g$, we can compute the feature vector at the bootstrapped return $r + \gamma g$ by a simple matrix multiplication.

What this enables: if this relationship holds exactly, then the mean embeddings of the true return distributions $U^\pi(x) = \mathbb{E}^\pi_x[\phi(\sum_{t=0}^\infty \gamma^t R_t)]$ satisfy their own exact linear Bellman equation:

$U^\pi(x) = \mathbb{E}^\pi_x[B_R U^\pi(X')]$

The derivation (Equation 4) proceeds in three steps: (a) from the distributional Bellman equation, $\mathbb{E}^\pi_x[\phi(R + \gamma G^\pi(X'))] = \mathbb{E}^\pi_x[\phi(\sum_{t=0}^\infty \gamma^t R_t)]$, where $G^\pi(x)$ are random variables distributed as the true return from state $x$; (b) applying the closure property, $\phi(R + \gamma G^\pi(X')) = B_R \phi(G^\pi(X'))$; (c) exchanging the linear map $B_R$ and the conditional expectation, $\mathbb{E}^\pi_x[B_R \phi(G^\pi(X'))] = \mathbb{E}^\pi_x[B_R \mathbb{E}[\phi(G^\pi(X')) \mid X']] = \mathbb{E}^\pi_x[B_R U^\pi(X')]$, which relies crucially on $B_R$ being linear.

Why this is restrictive (Bellman closure characterization). Rowland et al. (2019, Theorem 4.3) proved that the only mean embedding sketches satisfying this exact closure property are those whose feature functions span the same space as the first $m$ monomials $\{1, g, g^2, \ldots, g^{m-1}\}$. Equivalently, the sketch must be an invertible linear combination of the first $m$ moments of the return. This recovers Sobel's (1982) moment dynamic programming but inherits a critical practical problem: the $k$-th moment scales as $O((\text{return range})^k)$, so for $m = 50$, the moment magnitudes span dozens of orders of magnitude, making a single learning rate impossible to tune for TD learning (Appendix D.1, Figure 7). The polynomial feature space is also global — each feature responds everywhere on the real line — making it difficult to decode localized distributional information.

The least-squares relaxation. To escape the Bellman closure limitation while preserving the linear structure, the paper defines the Bellman coefficients $B_r$ as the solution to a least-squares regression problem (Equation 5):

$B_r := \arg\min_{B \in \mathbb{R}^{m \times m}} \mathbb{E}_{G \sim \mu}\left[\|\phi(r + \gamma G) - B \phi(G)\|_2^2\right]$

where $\mu$ is a weighting distribution over possible returns that determines which $g$ values the regression prioritizes fitting well, and $\|\cdot\|_2$ is the Euclidean norm in $\mathbb{R}^m$.

What it computes: for each reward $r$, the regression finds the matrix $B_r$ that best predicts $\phi(r + \gamma g)$ from $\phi(g)$ on average over returns $g$ drawn from $\mu$. Crucially, $B_r$ does not depend on $g$ — it is a single $m \times m$ matrix — but the fit is approximate: $\phi(r + \gamma g) \approx B_r \phi(g)$ with some residual error that depends on $g$. The error is zero only when $\phi$ is Bellman-closed (moments), and is nonzero but controllable for general $\phi$.

Why regression and not a nonlinear predictor: Remark 3.2 addresses the natural question of why not fit a nonlinear function $H(r, \phi(g))$ to get better accuracy. If $H$ is nonlinear in its second argument, then $\mathbb{E}[H(r, \phi(G(X')))] \neq H(r, \mathbb{E}[\phi(G(X'))])$, and the crucial step of exchanging expectation and prediction fails — step (c) in the derivation above is invalid. Linearity of the predictor in $\phi(g)$ is what lets us pull $B_r$ outside the conditional expectation and obtain the sketch Bellman equation $U^\pi(x) \approx \mathbb{E}^\pi_x[B_R U^\pi(X')]$. The approximation error from using a linear fit is the price paid for the ability to propagate sketch values in closed form.

Closed-form solution for $B_r$. Under the mild condition that the matrix $C = \mathbb{E}_{G \sim \mu}[\phi(G)\phi(G)^\top]$ is invertible (requiring that the features are not linearly dependent under $\mu$), the least-squares problem has the closed-form solution (Equation 7):

$B_r = C_r C^{-1}, \quad \text{where} \quad C = \mathbb{E}_{G \sim \mu}[\phi(G)\phi(G)^\top], \quad C_r = \mathbb{E}_{G \sim \mu}[\phi(r + \gamma G)\phi(G)^\top]$

What these matrices compute: $C \in \mathbb{R}^{m \times m}$ is the uncentered second-moment matrix of the feature vectors under the weighting distribution $\mu$ — each entry $C_{ij} = \mathbb{E}_{G \sim \mu}[\phi_i(G)\phi_j(G)]$ measures the correlation between features $i$ and $j$ over the return range. $C_r \in \mathbb{R}^{m \times m}$ is the uncentered cross-moment matrix between the shifted features $\phi(r + \gamma G)$ and the original features $\phi(G)$ — each entry $(C_r)_{ij} = \mathbb{E}_{G \sim \mu}[\phi_i(r + \gamma G)\phi_j(G)]$ measures how feature $j$ at the original return $G$ predicts feature $i$ at the bootstrap return $r + \gamma G$. The product $C_r C^{-1}$ is the standard multivariate linear regression coefficient matrix: it "projects" the target $\phi(r + \gamma G)$ onto the predictor $\phi(G)$ in the $L^2(\mu)$ sense.

Why this form: this is the vector-valued analog of scalar linear regression coefficients $b = \text{Cov}(X, Y)/\text{Var}(X)$. The $m^2$ entries of $B_r$ capture all cross-predictive relationships between the $m$ original features and the $m$ shifted features. The $O(m^3)$ cost of inverting $C$ is paid once during precomputation, after which each Bellman backup costs only $O(m^2)$ for the matrix-vector product $B_r U(x')$.

Choice of the weighting distribution $\mu$. The regression distribution $\mu$ determines which return values $g$ the Bellman coefficients optimize for. The paper uses a uniform distribution over a fine grid spanning the estimated return range $[\hat{G}_{\min} - b\hat{L}, \hat{G}_{\max} + b\hat{L}]$, where $b \approx 0.2$ provides a small buffer to ensure coverage of the full support (Appendix B.3). The grid is "densely spaced" — in the deep RL experiments, 100,000 points evenly spaced over $[-10, 10]$ (Appendix C.2). This choice ensures that the regression error is roughly uniform over the return range, rather than concentrated in low-probability regions. For some feature-map/$\mu$ combinations, the integrals defining $C$ and $C_r$ can be computed analytically — the paper provides explicit formulas for Gaussian $\phi_i(z) = \exp(-s^2(z - z_i)^2/2)$ with $\mu$ as Lebesgue measure (Appendix B.3), where:

$C_{ij} = \sqrt{\frac{\pi}{2}} s \exp\left(-\frac{s^2(z_i - z_j)^2}{2}\right)$

$(C_r)_{ij} = \sqrt{\frac{\pi}{s(1 + \gamma^2)}} \exp\left(-\frac{s(r + \gamma z_i - \gamma z_j)^2}{1 + \gamma^2}\right)$

Regularization in practice. The paper adds a small $L^2$ regularizer with weight $10^{-9}$ when solving for $B_r$ in the deep RL experiments to avoid numerical instability from near-singular $C$ (Appendix C.2), tuned from $\{10^{-15}, 10^{-12}, 10^{-9}, 10^{-6}, 10^{-3}\}$. The solution becomes $B_r = C_r(C + \lambda I)^{-1}$ where $\lambda$ is the regularization strength.

Computational properties of $B_r$. The paper discusses structure that can be exploited (Appendix B.2). For "binning features" (indicator functions as in Proposition 4.4), the features have disjoint or nearly-disjoint support, so $\phi_i(G)\phi_j(G) \approx 0$ for $i \neq j$, making $C$ diagonally dominant and $B_r$ a narrow-band matrix — the matrix-vector product becomes $O(m)$ rather than $O(m^2)$. Similar sparsity holds approximately for localized features like low-bandwidth Gaussians, and the paper suggests truncating near-zero coefficients to exploit this.

Handling unknown or infinite reward sets. For TD learning where the set of possible rewards $\mathcal{R}$ may not be known in advance, the paper describes two strategies (Appendix B.4): (1) precompute $C^{-1}$ and compute $C_r$ online upon observing a new reward $r$, reducing marginal cost to one matrix-matrix product $C_r C^{-1}$; (2) learn an approximator $H: \mathbb{R} \to \mathbb{R}^{m \times m}$ mapping rewards to Bellman coefficients, and use its predictions as proxies to avoid solving the regression for every new reward value.

Empirical properties and non-normal dynamics. The paper analyzes $B_r$ for sigmoid and Gaussian base features with $\gamma = 0.8$, $r = 1$, and $m = 20$ evenly spaced anchors in $[-8, 8]$ (Appendix B.5, Figure 6). Key observations: (1) the maximum absolute difference between $\phi(r + \gamma z)$ and $B_r \phi(z)$ over the regression grid is less than $0.002$, indicating near-perfect fit within the training range; (2) the largest singular values of $B_r$ can exceed 1, meaning a single application of $B_r$ may expand the input norm — $B_r$ is not a contraction in $L^2$; (3) however, all eigenvalues have real parts less than or very close to 1, suggesting that repeated multiplication $B_r^k u$ converges to a stable fixed point despite transient expansion. This non-normal dynamics — where a matrix is not contractive in Euclidean norm but is stable under iteration — is highlighted as potentially relevant for biological implementations of distributional RL (Hennequin et al., 2012; Bondanelli & Ostojic, 2020).

The constant feature exception. When $\phi_1(g) \equiv 1$ (the constant feature), the sketch Bellman operator becomes affine rather than linear — the first coordinate of $U(x)$ is constrained to equal 1, and the first row of $B_r$ is $(1, 0, \ldots, 0)$ with $r$ in the appropriate position to recover the standard expected-value Bellman equation. The paper notes this is "very crucial for a good performance" (Appendix C.2) in deep RL, and includes it by swapping to an appended constant feature rather than learning it.

3.4.3 From Bellman Coefficients to the Sketch Bellman Operator

The sketch Bellman operator $\mathcal{T}^\pi_\phi$. Given precomputed Bellman coefficients $B_r$ for each $r \in \mathcal{R}$, the sketch Bellman operator $\mathcal{T}^\pi_\phi: (\mathbb{R}^m)^\mathcal{X} \to (\mathbb{R}^m)^\mathcal{X}$ is defined for any collection of sketch vectors $U: \mathcal{X} \to \mathbb{R}^m$ by:

$(\mathcal{T}^\pi_\phi U)(x) := \mathbb{E}^\pi_x[B_R U(X')]$

where the expectation is over the random transition $(X = x, A \sim \pi(\cdot|x), R \sim P_R(x, A), X' \sim P(\cdot|x, A))$, and $B_R$ denotes the Bellman coefficient matrix corresponding to the sampled reward $R$.

What it computes: at each state $x$, the operator computes a new sketch vector by (1) for each possible transition to a next state $x'$ with reward $r$, looking up the sketch vector $U(x')$ at the successor state, (2) applying the Bellman coefficient matrix $B_r$ to produce $B_r U(x')$ — the predicted sketch value of the bootstrap return distribution — and (3) averaging these predictions over all possible transitions weighted by transition probabilities and policy probabilities. In the common case of conditional independence $R \perp\!\!\!\perp X' \mid X$ (rewards depend only on state, not jointly on state and next state), this simplifies to:

$U(x) \leftarrow \mathbb{E}^\pi_x[B_R] \mathbb{E}^\pi_x[U(X')] = \mathbb{E}^\pi_x[B_R] \sum_{x' \in \mathcal{X}} P(x' \mid x) U(x')$

where $\mathbb{E}^\pi_x[B_R]$ is the expected Bellman coefficient under the reward distribution at state $x$ — a single $m \times m$ matrix per state that can be precomputed.

Why this operator replaces SFDP's imputation round-trip. In SFDP (Rowland et al., 2019), a Bellman backup requires: (i) $\iota(U)$ — impute an approximate distribution from sketch values; (ii) $\mathcal{T}^\pi(\iota(U))$ — apply the full distributional Bellman operator (shifting, scaling, and mixing distributions); (iii) $\psi(\mathcal{T}^\pi(\iota(U)))$ — extract new sketch values. Each step is computationally intensive and step (i) requires a sketch-specific inverse map. The sketch Bellman operator collapses all three steps into a single linear operation: $U \leftarrow \mathbb{E}^\pi_x[B_R U(X')]$. The imputation and distributional backup are implicitly performed by the Bellman coefficients $B_r$, which were optimized offline to approximate the composition $\phi \circ (r + \gamma \cdot) \circ \psi^{-1}$.

Why $\mathcal{T}^\pi_\phi$ is linear (and why that matters). The operator acts on the vector space $(\mathbb{R}^m)^\mathcal{X}$ of all sketch-value assignments, and satisfies $\mathcal{T}^\pi_\phi(\alpha U_1 + \beta U_2) = \alpha \mathcal{T}^\pi_\phi(U_1) + \beta \mathcal{T}^\pi_\phi(U_2)$ for any scalars $\alpha, \beta$, because $B_r$ is a matrix multiplication and expectation is linear. This linearity is what enables the clean error propagation analysis in Section 4 — the operator's Lipschitz constant with respect to any norm is simply the operator norm of the transition-averaged Bellman coefficient matrix, and the asymptotic behavior is governed by linear fixed-point iteration. In contrast, SFDP operators are nonlinear due to the imputation step and distributional projection, making analysis significantly more complex.

Non-contractivity. The paper does not claim that $\mathcal{T}^\pi_\phi$ is a contraction in any norm. In fact, the singular value analysis of $B_r$ (Appendix B.5) shows that $\|B_r\|_2$ can exceed 1, meaning the operator may expand differences between sketch vectors rather than shrink them — in contrast to the standard expected-value Bellman operator, which is a $\gamma$-contraction in supremum norm. The convergence of Sketch-DP iterates therefore relies not on contractivity but on the stability of the linear dynamical system $U_{k+1} = \mathcal{T}^\pi_\phi U_k$ — the eigenvalues of the transition-averaged operator having magnitude less than or equal to 1, ensuring bounded trajectories — combined with the propagation analysis showing convergence to a neighborhood of the true sketch values whose radius is proportional to the regression error.

3.4.4 Sketch Dynamic Programming (Sketch-DP)

Algorithm structure. Sketch-DP (Algorithm 1, DP branch) is a direct iterative method for policy evaluation given a known MDP. The algorithm has two phases:

Phase 1: Precomputation. Before any DP iterations, the Bellman coefficients are computed once:

1. Choose a weighting distribution $\mu$ (typically a uniform fine grid spanning the estimated return range with a small buffer).
2. Compute the second-moment matrix $C = \mathbb{E}_{G \sim \mu}[\phi(G)\phi(G)^\top] \in \mathbb{R}^{m \times m}$.
3. For each reward $r \in \mathcal{R}$, compute the cross-moment matrix $C_r = \mathbb{E}_{G \sim \mu}[\phi(r + \gamma G)\phi(G)^\top] \in \mathbb{R}^{m \times m}$.
4. For each $r$, compute the Bellman coefficient matrix $B_r = C_r C^{-1}$ (with optional $L^2$ regularization as $C_r(C + \lambda I)^{-1}$).

For environments where rewards and next states are conditionally independent given the current state, the per-state expected Bellman coefficient $\bar{B}(x) = \mathbb{E}^\pi_x[B_R] = \sum_{r} P_R(r \mid x) B_r$ can also be precomputed for faster iteration.

Phase 2: Iterative updates. Starting from an initial sketch estimate $U_0: \mathcal{X} \to \mathbb{R}^m$ (typically $U_0(x) = \phi(0)$ for all states, representing the sketch of a Dirac delta at zero return), the DP iteration proceeds:

$U_{k+1}(x) \leftarrow \sum_{x', r} P(r, x' \mid x, \pi) B_r U_k(x') \quad \text{for all } x \in \mathcal{X}$

where $P(r, x' \mid x, \pi) = \sum_a \pi(a \mid x) P(x' \mid x, a) P_R(r \mid x, a)$ if not pre-aggregated. In the conditionally independent case: $U_{k+1}(x) \leftarrow \bar{B}(x) \sum_{x'} P(x' \mid x, \pi) U_k(x')$. The iteration terminates after a fixed number of steps or when changes fall below a threshold.

What each iteration computes: at each state $x$, the update takes the current sketch estimates at all possible successor states $x'$, weights them by transition probabilities, and then applies the state- and reward-averaged Bellman coefficient matrix — effectively performing a linear Bellman backup entirely in $\mathbb{R}^m$. There is no decoding to probability distributions, no application of the distributional Bellman operator on infinite-dimensional objects, and no re-encoding. The entire backup is a single matrix-vector product (or a sum of matrix-vector products if rewards are not pre-aggregated).

Computational complexity. The precomputation phase costs $O(m^3)$ for the matrix inversion of $C$ plus $O(|\mathcal{R}| m^3)$ if computing $C_r C^{-1}$ by matrix multiplication (or $O(|\mathcal{R}| m^2)$ if solving linear systems). Each DP iteration costs $O(|\mathcal{X}|^2 m^2)$ in the worst case (full transition matrix), or $O(|\mathcal{X}| \cdot \text{avg-degree} \cdot m^2)$ for sparse transition graphs, since each state-to-state transition requires an $m \times m$ matrix-vector product. This is substantially cheaper than SFDP, whose per-iteration cost includes solving a quadratic program per state. The paper's empirical comparison (Appendix D.4, Figure 11) shows Sketch-DP being over 100× faster per iteration than SFDP for $m = 75$ features/expectiles.

Visualization via imputation. To produce interpretable results, the paper optionally decodes the learned sketch values back into probability distributions using an imputation strategy $\iota: \mathbb{R}^m \to \mathscr{P}(\mathbb{R})$ (Appendix B.1). The imputation solves:

$\iota(u) = \arg\min_{p \in \Delta_n} \left\|\sum_{i=1}^n p_i \phi(z_i) - u\right\|_2^2$

where $\Delta_n$ is the probability simplex over $n$ support points $\{z_1, \ldots, z_n\}$ (typically the feature anchors, with jitter to avoid alignment artifacts). This is a convex quadratic program solved efficiently via SciPy's MINIMIZE. The imputation is used only for evaluation and visualization — it is not part of the DP update and does not affect the learned sketch values.

3.4.5 Sketch Temporal-Difference Learning (Sketch-TD)

Algorithm structure. Sketch-TD (Algorithm 1, TD branch) adapts the sketch Bellman operator to the sample-based, incremental setting of reinforcement learning where transition dynamics are unknown. Given a stream of transitions $(x, a, r, x')$ generated by following policy $\pi$, the update for the sketch estimate at state $x$ is:

$U(x) \leftarrow (1 - \alpha) U(x) + \alpha B_r U(x')$

where $\alpha \in (0, 1)$ is a learning rate.

What it computes: this is a standard stochastic approximation update — it moves the current estimate $U(x)$ a small step $\alpha$ toward the bootstrap target $B_r U(x')$, which is the sketch of the distribution of $r + \gamma G^\pi(x')$ predicted from the successor state's sketch. The Bellman coefficient $B_r$ is looked up based on the observed reward $r$ (either from a precomputed table or computed online if $\mathcal{R}$ is unknown a priori). The update does not require knowing the transition probabilities — it uses the observed sample transition as an unbiased estimate of the expectation in the DP update, assuming that $(x, r, x')$ is drawn from the correct transition distribution.

Why this form: it is the natural sample-based analog of Sketch-DP, exactly as TD(0) is the sample-based analog of iterative policy evaluation in expected-value RL. The key substitution is replacing the expectation $\mathbb{E}^\pi_x[B_R U(X')]$ with the single-sample estimate $B_r U(x')$ and using an exponential moving average with learning rate $\alpha$ to handle stochasticity. The linearity of the target in $U(x')$ means this is a linear stochastic approximation algorithm — it fits into the standard Robbins-Monro framework (Kushner & Yin, 1997; Bertsekas & Tsitsiklis, 1996), though the paper does not provide a full convergence proof for the TD case (noting it as future work).

The difficulty with polynomial features. Figure 7 (Appendix D.1) empirically demonstrates why moment-based sketches (polynomial features) fail for TD learning at scale. With $m = 50$ polynomial features $\phi_i(g) = g^{i-1}$, the sketch components span enormously different scales — the mean is $O(1)$, while the 50th moment can be $O(10^{50})$. A single learning rate $\alpha$ cannot simultaneously be small enough to prevent divergence in high-order components and large enough to make meaningful progress in low-order components. Sigmoid-based sketches of equal dimensionality ($m = 50$) show a wide basin of good learning rates (roughly $10^{-4}$ to $10^{-2}$ depending on the environment) with substantially lower Cramér distances to ground truth, because all features are bounded in $[0, 1]$ and vary over comparable ranges.

Online computation of $B_r$ for unknown rewards. When the set of possible rewards $\mathcal{R}$ is unknown or infinite, Algorithm 1 is modified to compute $B_r$ on-the-fly: precompute and cache $C^{-1}$ once, then when a new reward $r$ is observed for the first time, compute $C_r$ (requiring evaluating the integral $\mathbb{E}_{G \sim \mu}[\phi(r + \gamma G)\phi(G)^\top]$, e.g., via numerical integration or analytic formulas) and set $B_r = C_r C^{-1}$. The marginal cost is $O(m^3)$ for the matrix multiplication (or $O(m^2)$ if $C_r$ has exploitable structure), amortized over all future occurrences of reward $r$. For settings where this is still too expensive, the paper suggests learning a function $H: \mathbb{R} \to \mathbb{R}^{m \times m}$ that maps reward values to Bellman coefficients (Remark 3.2, Appendix B.4), though this is not implemented in the paper's experiments.

Synchronous vs. asynchronous updates. In the tabular experiments (Appendix D.1), Sketch-TD is run with synchronous updates — all states are updated simultaneously using the previous iteration's sketch values — to provide a controlled comparison between TD and DP. This is a stronger condition than typical asynchronous TD (where updates use whatever values are currently stored) and makes the TD dynamics more directly comparable to the DP iteration.

3.4.6 Convergence Analysis of Sketch-DP

The error propagation framework. The convergence analysis (Section 4) addresses the fundamental question: since $\mathcal{T}^\pi_\phi$ is not exact (the Bellman coefficients introduce regression error) and the sketch $\Phi$ is lossy (distributions with different sketch values may be close, and vice versa), how far can the Sketch-DP iterates $U_k$ drift from the true sketch values $U^\pi$? The analysis decomposes the error into three components, illustrated in Figure 3.

Component 1: Single-step regression error (Proposition 4.1). Let $\Phi(\nu) = \mathbb{E}_{Z \sim \nu}[\phi(Z)]$ be the sketch operator from distributions to $\mathbb{R}^m$, and let $\Phi$ applied to a return-distribution function $\eta$ mean $(\Phi\eta)(x) = \Phi(\eta(x))$. The error from one approximate application of $\mathcal{T}^\pi_\phi$ instead of the exact $\Phi \circ \mathcal{T}^\pi$ is:

$\max_{x \in \mathcal{X}} \left\|\Phi(\mathcal{T}^\pi\eta)(x) - (\mathcal{T}^\pi_\phi \Phi\eta)(x)\right\| \leq \sup_{g \in [G_{\min}, G_{\max}]} \max_{r \in \mathcal{R}} \left\|\phi(r + \gamma g) - B_r \phi(g)\right\|$

where $\|\cdot\|$ is any norm on $\mathbb{R}^m$.

What it bounds: the left-hand side is the difference, at each state, between (i) taking the true distributional Bellman backup $\mathcal{T}^\pi\eta$ and then extracting its sketch, and (ii) taking the sketch of $\eta$ and applying the approximate sketch operator $\mathcal{T}^\pi_\phi$. The right-hand side is the worst-case pointwise regression error — the maximum over all returns $g$ and rewards $r$ of how poorly $B_r\phi(g)$ approximates $\phi(r + \gamma g)$. This is the Bellman approximation error $\varepsilon_B$ in Proposition 4.2.

Why this form: the proof (Appendix A) uses the linearity of expectation to pull the difference inside and apply Jensen's inequality. The supremum over $g$ is inside the maximum over $r$, meaning the worst-case reward-regression pair determines the per-step error bound. This bound is deterministic and data-independent given $\phi$, $B_r$, and the return range — it can be evaluated numerically before running the algorithm to assess the quality of the Bellman coefficients.

Component 2: Reconstruction error (embedding to distribution). For any two return-distribution functions $\eta, \bar{\eta}$ with sketches $U, \bar{U}$, there exists a metric $d$ on distributions (the paper uses supremum-Wasserstein or supremum-Cramér distance) and a constant $\varepsilon_R \geq 0$ such that:

$d(\eta, \bar{\eta}) \leq \|U - \bar{U}\|_\infty + \varepsilon_R$

What it says: two return-distribution functions whose sketches are close in supremum norm (over states) have close true distributions, up to an additive slack $\varepsilon_R$. This slack captures the reconstruction error — information lost by the sketch. If the sketch is rich enough, $\varepsilon_R$ is small; for a degenerate sketch that maps all distributions to the same vector, $\varepsilon_R$ would be large.

Component 3: Embedding error (distribution to embedding). The reverse direction:

$\|U' - \bar{U}'\|_\infty \leq d(\eta', \bar{\eta}') + \varepsilon_E$

What it says: the sketch vectors of two return-distribution functions cannot be further apart (in supremum norm) than the distance between the true distributions plus a slack $\varepsilon_E$ that captures non-injectivity — distributions that are far apart in $d$ may have identical sketches, and $\varepsilon_E$ bounds how much closer the sketches can be than the distributions.

The error propagation lemma (Proposition 4.2). Under the three error bounds above, and assuming $d$ is a metric under which $\mathcal{T}^\pi$ is a $\gamma_c$-contraction (the paper notes that both supremum-Wasserstein and supremum-Cramér distances satisfy this with $\gamma_c \leq \gamma$), for any $\eta, \bar{\eta}$ with sketches $U, \bar{U}$ satisfying $\|U - \bar{U}\|_\infty \leq \delta$, we have:

$\|\Phi\mathcal{T}^\pi\eta - \mathcal{T}^\pi_\phi \bar{U}\|_\infty \leq \gamma_c(\delta + \varepsilon_R) + \varepsilon_B + \varepsilon_E$

What it computes: the error after one combined step of (i) applying the true Bellman operator to $\eta$ and extracting the sketch, versus (ii) applying the approximate sketch operator to $\bar{U}$. The proof (Appendix A) traces a path through Figure 3: from $\bar{U}$ to $\bar{\eta}$ (reconstruction error $\delta + \varepsilon_R$), applying $\mathcal{T}^\pi$ (contracts by $\gamma_c$), extracting the sketch (embedding error $\varepsilon_E$), and comparing to the approximate operator on $\bar{U}$ (Bellman approximation error $\varepsilon_B$).

The main convergence result (Proposition 4.3). Applying the error propagation lemma to the sequence $U_{k+1} = \mathcal{T}^\pi_\phi U_k$ and $U^\pi$, and taking the limit supremum:

$\limsup_{k \to \infty} \|U_k - U^\pi\|_\infty \leq \frac{1}{1 - \gamma_c}(\gamma_c \varepsilon_R + \varepsilon_B + \varepsilon_E)$

What it says: the asymptotic error is bounded by a linear combination of the three error sources, amplified by the horizon factor $1/(1-\gamma_c)$. Critically, all three terms ($\varepsilon_B, \varepsilon_R, \varepsilon_E$) can be made arbitrarily small by increasing the number of features $m$ for well-designed sketches, meaning the algorithm can achieve arbitrary accuracy as $m \to \infty$. The bound assumes that $\mathcal{T}^\pi$ maps distributions supported on $[G_{\min}, G_{\max}]$ to themselves and that $\mathcal{T}^\pi_\phi$ maps the corresponding sketch set to itself — conditions that hold for the indicator features analyzed in Proposition 4.4.

Concrete instantiation for indicator features (Proposition 4.4). To make the abstract theory concrete, the paper fully works out the bounds for a specific sketch: indicator functions $\phi_i(z) = \mathbf{1}\{z_i \leq z < z_{i+1}\}$ for $i = 1, \ldots, m-1$, with $\phi_m(z) = \mathbf{1}\{z_1 \leq z \leq z_{m+1}\}$, where $\{z_1, \ldots, z_{m+1}\}$ is an equally-spaced grid over $[G_{\min}, G_{\max}]$ of width $\Delta = (G_{\max} - G_{\min})/m$, and $\mu = \text{Unif}([G_{\min}, G_{\max}])$. Under the norm $\|u\| = \Delta \sum_{i=1}^m |u_i|$:

$\limsup_{k \to \infty} \|U_k - U^\pi\|_\infty \leq \frac{(G_{\max} - G_{\min})(3 + 2\gamma)}{(1 - \gamma)m}$

How the bound is derived (Appendix A):

• Reconstruction error $\varepsilon_R = 2\Delta$: for any distribution $\nu$ supported on $[z_1, z_{m+1}]$, define its projection $\Pi\nu$ by mapping each point to the greatest grid point below it. The 1-Wasserstein distance $w_1(\nu, \Pi\nu) \leq \Delta$ since mass moves at most $\Delta$. For two projected distributions, $w_1(\Pi\nu, \Pi\nu') = \|\Phi\nu - \Phi\nu'\|$ (the norm is exactly the $L^1$ distance between their sketch vectors). Chaining via triangle inequality gives $w_1(\nu, \nu') \leq \|\Phi\nu - \Phi\nu'\| + 2\Delta$, and similarly for the reverse direction with $\varepsilon_E = 2\Delta$.
• Bellman approximation error $\varepsilon_B = \Delta$: because $r + \gamma G$ varies over an interval of width $\gamma\Delta$ when $G$ is uniform over $[z_i, z_{i+1})$, $\phi(r + \gamma G)$ takes at most two distinct values, and the regression optimal $B_r$ can fit this with maximum per-coordinate error bounded by $\Delta$ (in the chosen norm).
• Plugging into the general bound yields the $O(1/m)$ result.

What this means: with $m$ indicator features, Sketch-DP converges to within $O(1/m)$ of the true sketch values, and this error can be driven arbitrarily close to zero by increasing $m$. The $O(1/m)$ rate is the natural parametric rate for approximating a distribution with histogram bins of width $O(1/m)$. This is the key theoretical justification for the framework: approximate closure is sufficient for principled convergence, and the approximation error can be systematically controlled.

Interpretation and limitations. The theory proves convergence to a neighborhood, not exact convergence — the error floor scales as $1/m$. It does not prove contractivity of $\mathcal{T}^\pi_\phi$ (which, as discussed, may expand in some directions). Instead, it uses the contraction of the true distributional operator $\mathcal{T}^\pi$ in Wasserstein/Cramér distance to bound how errors propagate through the composition, showing that the approximate operator tracks the true operator with bounded per-step error that compounds sub-geometrically. The analysis is specific to the DP case; convergence of Sketch-TD is not analyzed and is noted as future work.

Why this analysis is novel. Prior work on distributional RL convergence (Rowland et al., 2018; Dabney et al., 2018b; Rowland et al., 2023) analyzed specific distribution representations (categorical, quantile) by proving that the algorithmic operator is a contraction in some metric. This paper's analysis is fundamentally different: $\mathcal{T}^\pi_\phi$ is not proved contractive; instead, it is shown to be an approximate simulator of the true sketch operator $\Phi \circ \mathcal{T}^\pi$, with bounded per-step error. This error propagation approach (inspired by Munos, 2003; Wu et al., 2023) is more general and handles the approximate nature of the Bellman coefficients naturally.

3.4.7 Sketch-DQN: Adaptation to Deep Reinforcement Learning

Architecture. The deep RL agent (Section 5.1, Appendix C.2) parameterizes sketch values as a neural network $U_\theta: \mathcal{X} \times \mathcal{A} \to \mathbb{R}^m$ that maps state-action pairs to $m$-dimensional sketch vectors. The architecture follows QR-DQN (Dabney et al., 2018b): a convolutional torso (same as DQN) feeds into a fully-connected layer that outputs $m$ values per action — these are the predicted sketch coordinates $\hat{U}_\theta(x, a)_i$ for $i = 1, \ldots, m$. Key modifications from QR-DQN:

1. Output nonlinearity: a sigmoid or tanh nonlinearity is applied to the final layer output to bound the predicted sketch values, matching the known output range of the chosen base feature $\kappa$ (e.g., sigmoid features are in $[0, 1]$, so sigmoid output nonlinearity is natural). This improves training stability.

2. Constant feature handling: the network predicts only the $m-1$ non-constant dimensions of the sketch; the constant feature $\phi_1(g) \equiv 1$ is appended as a hard-coded value after the network output, ensuring the sketch operator remains affine rather than purely linear. This was found to be "very crucial for a good performance" (Appendix C.2).

3. Value readout: to define a greedy policy (required for Q-learning), expected returns must be read out from sketch values. The paper precomputes value-readout coefficients $\beta \in \mathbb{R}^m$ by solving the least-squares problem:

$\beta := \arg\min_{\beta \in \mathbb{R}^m} \mathbb{E}_{G \sim \mu}\left[(G - \langle \beta, \phi(G) \rangle)^2\right]$

What this computes: the linear combination $\langle \beta, \phi(G) \rangle = \sum_{i=1}^m \beta_i \phi_i(G)$ is the best linear predictor of the return $G$ from the feature vector $\phi(G)$, in the $L^2(\mu)$ sense. Given a predicted sketch vector $\hat{U}_\theta(x, a)$ (which estimates $\mathbb{E}[\phi(G^\pi(x, a))]$), the scalar $\langle \beta, \hat{U}_\theta(x, a) \rangle$ estimates the expected return $\mathbb{E}[G^\pi(x, a)]$. This is used to select the greedy action and to compute the Q-learning target's action-value.

Why this form: solving for $\beta$ is again a least-squares problem with closed-form solution $\beta = C^{-1} \mathbb{E}_{G \sim \mu}[G \phi(G)]$, using the same $C$ matrix computed for the Bellman coefficients. The readout is linear in the sketch values, preserving the overall linearity of the framework. The $\beta$ vector effectively tells us how to weight the different basis function expectations to reconstruct the mean return — it is the mean embedding analog of the inverse-link function in generalized linear models.

Q-learning-style update rule. Given a transition $(x, a, r, x')$, Sketch-DQN computes:

$a' = \arg\max_{\tilde{a}} \langle \beta, U_{\bar{\theta}}(x', \tilde{a}) \rangle$

using the target network parameters $\bar{\theta}$ to select the best next action, and then updates the online network parameters $\theta$ by minimizing:

$\nabla_\theta \left\|U_\theta(x, a) - B_r U_{\bar{\theta}}(x', a')\right\|_2^2$

What it computes: the target for the sketch vector at $(x, a)$ is the Bellman coefficient $B_r$ applied to the target network's sketch prediction at the greedily-selected next state-action pair $(x', a')$. The loss is the squared Euclidean distance between the online network's prediction and this target, matching the standard DQN template but operating on $m$-dimensional vectors rather than scalars. The gradient is taken only with respect to $\theta$ (the target network $\bar{\theta}$ is held fixed during the gradient step).

Why this form: the target $B_r U_{\bar{\theta}}(x', a')$ is a constant vector for the gradient computation — it is treated as a supervised learning target, not differentiated through. This is the standard semi-gradient approach in TD learning: the bootstrap target is computed with the target network to stabilize training, and only the prediction head is updated. The vector-valued loss means all $m$ sketch coordinates are trained jointly, sharing the same convolutional features.

Feature map and hyperparameters for Atari. The results in Figure 5 used:

• Base feature: sigmoid $\kappa(x) = 1/(1 + \exp(-x))$
• Number of features: $m = 401$ (tuned from $\{101, 201, 401\}$), plus one constant feature appended
• Slope: $s = 5$ (tuned from $\{1, 2, \ldots, 12\}$; larger slopes caused worst-case regression error exceeding $0.01$ and were rejected)
• Anchors: 401 evenly spaced points in $[-12, 12]$ (loosely motivated by the C51 atom range, Bellemare et al., 2017)
• Regression distribution $\mu$: 100,000 uniform points in $[-10, 10]$, with $L^2$ regularization $10^{-9}$
• Learning rate: $5 \times 10^{-5}$ (same as QR-DQN default; found optimal among sweeps)
• Training frames: 200 million (standard Atari benchmark)
• Other hyperparameters (exploration, replay buffer, target network update frequency): identical to QR-DQN

The paper also tested Gaussian base features with $m = 201$, slope $s = 1.67$, and anchors in $[-12, 12]$ (Figure 13), but found sigmoid features more performant. The tuning process for the slope involved rejecting hyperparameters where the worst-case regression error $\max_{r \in \{-1, 0, 1\}} \max_{g \in \text{supp}(\mu)} \|\phi(r + \gamma g) - B_r \phi(g)\|$ exceeded $0.01$ — a practical screen based on Proposition 4.1's insight that the per-step error is bounded by this quantity.

Runtime comparison. Table 1 (Appendix D.6) reports training frame rates on a single V100 GPU:

• Sketch-DQN ($m = 401$): $1326 \pm 107$ frames/second
• C51 (51 atoms): $1309 \pm 146$ fps
• QR-DQN (201 quantiles): $1258 \pm 107$ fps
• IQN (64 quantiles): $1120 \pm 90$ fps; IQN (201 quantiles): $698 \pm 41$ fps

Sketch-DQN is the fastest among the distributional methods, with IQN being dramatically slower at high quantile counts because its architecture requires a separate forward pass through the MLP component for each quantile level. The sketch architecture, like QR-DQN, produces all $m$ predictions from a single final hidden layer, keeping the per-step cost close to that of DQN. The $O(m^2)$ Bellman coefficient multiplication is an additional cost but is precomputed and cached — at runtime, the target computation $B_r U_{\bar{\theta}}(x', a')$ is a single matrix-vector product costing $O(m^2) = 401^2 \approx 160{,}000$ multiply-adds, negligible relative to the convolutional forward pass.

3.4.8 Design Choices and Their Justifications

Why the feature map $\phi$ is user-chosen rather than learned. The entire framework is parameterized by $\phi$ — change $\phi$ and you get a different family of distributional RL algorithms. The paper treats $\phi$ as a design choice rather than a learned component to maintain the linear structure (if $\phi$ were learned online, the Bellman coefficients would need to be continuously recomputed) and to provide interpretable control over which distributional features are captured. The translation family (Equation 8) provides a systematic way to explore this design space, with the three knobs — base function $\kappa$, count $m$, slope $s$ — controlling the expressiveness-resolution tradeoff in an intuitive way (more features = finer representation, wider slope = smoother representation).

Why Bellman coefficients are precomputed rather than learned online. Precomputing $B_r$ from a fixed weighting distribution $\mu$ decouples the coefficient optimization from the RL algorithm and ensures the regression error is controlled before training begins. Learning the coefficients online would introduce non-stationarity (the distribution of returns changes as the policy evolves) and complicate the theoretical analysis. The price is that the coefficients are optimized for the return range specified by $\mu$, which must be chosen to cover the true returns — a reasonable assumption when return limits are known or can be bounded (e.g., $r_{\min}/(1-\gamma)$ to $r_{\max}/(1-\gamma)$).

Why the constant feature is appended rather than learned. Appending a hard-coded constant feature ensures that the first coordinate of $U(x)$ is always 1, which means the sketch operator is affine (recovers expected-value Bellman equation) rather than linear. If the constant feature were learned, the sketch operator would be purely linear, and the trivial solution $U(x) = 0$ for all $x$ would be a valid fixed point (since $B_r \cdot 0 = 0$). The constant feature breaks this degeneracy and anchors the representation.

Why the value readout $\beta$ is a separate linear readout rather than end-to-end Q-learning. The sketch values are learned via the sketch Bellman target $B_r U(x')$, not via direct regression on returns. To use the sketch for action selection, a separate mapping from sketch values to expected returns is needed. The linear readout $\beta$ provides this in a principled way (optimal linear predictor under $\mu$) without distorting the sketch learning objective — the sketch is still trained to match the bootstrap target, not to directly predict returns. This modularization separates the problems of learning distributional structure (what $\phi$ captures) from mean prediction (what $\beta$ reads out).

4. Key Insights and Innovations

Innovation 1: The Sketch Bellman Operator — Distributional RL Without Distributions

The paper's most fundamental conceptual move is the demonstration that distributional reinforcement learning can be performed entirely in the space of finite-dimensional sketch vectors, without ever explicitly constructing, manipulating, or sampling from probability distributions during Bellman backups. This is not merely an algorithmic optimization — it is a reframing of what distributional RL means computationally.

What the field did before. Prior work on sketch-based distributional RL (Rowland et al., 2019; Bellemare et al., 2023) treated sketches as a compression format for probability distributions: one stored sketch values for efficiency, but every Bellman backup required round-tripping through distribution space via an imputation strategy (ι), applying the full distributional Bellman operator to the imputed distribution (T^π), and re-extracting the sketch (ψ). This was computationally expensive, required designing sketch-specific pseudo-inverses, and made theoretical analysis cumbersome due to the composition of three nonlinear maps. Direct distributional methods (C51, QR-DQN, MMDRL) sidestep the imputation problem but at the cost of committing to specific distributional representations (categorical histograms, Dirac mixtures) rather than learning targeted statistical functionals. Both paradigms implicitly assume that distributional RL's primary object of computation is distributions themselves; sketches are either a storage format or not used at all.

What makes this distinctive. The insight is that the mean embedding sketch and the distributional Bellman operator can be composed approximately into a single linear operator on ℝ^m — the sketch Bellman operator T^π_φ — defined by precomputing Bellman coefficient matrices B_r that capture how the feature function φ transforms under the bootstrap map g ↦ r + γg. This operator is:

• Linear, making it analyzable with standard linear operator theory and enabling the error propagation framework of Section 4;
• Closed-form, requiring only matrix-vector products for each backup rather than solving optimization problems (imputation) or applying nonlinear particle updates (quantile regression, MMD gradients);
• Independent of any distribution representation — the algorithm never instantiates a probability distribution, never discretizes a CDF, never samples particles. It operates purely on the expectations of nonlinear basis functions.

The conceptual shift is from "distributions are the primary objects; sketches are a lossy encoding" to "sketches are the primary objects; distributions are implicit in the Bellman coefficients." The imputation step in SFDP is not accelerated or approximated — it is eliminated as a category. This is visible in Figure 1's visual comparison: the SFDP pathway has three stages with heavy computation, while Sketch-DP collapses them into a single linear update.

Why it matters beyond performance. This reframing has downstream consequences that the performance numbers alone don't capture:

• Biological plausibility. Tano et al. (2020) criticized imputation-based sketch methods as implausible models of neural distributional coding because neurons would need to explicitly reconstruct probability distributions and apply distributional Bellman backups. The sketch Bellman operator U(x) ← 𝔼[B_R U(X')] is a purely local, linear propagation rule — exactly the kind of computation that recurrent neural circuits could plausibly implement via synaptic weight matrices encoding B_r. The non-normal dynamics observed in Appendix B.5 (transient expansion followed by stable convergence) are also consistent with dynamical properties observed in neural circuits (Hennequin et al., 2012; Bondanelli & Ostojic, 2020). This is not a performance claim but an explanatory one: the framework provides a candidate computational mechanism for how brains might perform distributional RL without explicit probabilistic inference.

• Algorithmic generality. The framework provides a single recipe — choose φ, compute B_r via least squares — that works for any mean embedding sketch, including ones for which designing an imputation strategy would be difficult or unnatural (sinusoids, random Fourier features, learned basis functions). This dramatically lowers the barrier to exploring new sketches, since one no longer needs to solve a non-trivial inverse problem to operationalize each new functional.

• Theoretical tractability. Because T^π_φ is linear, the entire error propagation analysis (Propositions 4.1–4.3) reduces to bounding three quantities (ε_B, ε_R, ε_E) that have clean interpretations — regression error, reconstruction loss, embedding distortion — and then chaining them via the contraction of the true distributional operator. This is substantially cleaner than analyzing SFDP's ψ ◦ T^π ◦ ι composition, and it opens the door to importing tools from linear operator theory, spectral analysis, and perturbation theory into distributional RL analysis.

What kind of contribution this is. This is a fundamental reframing, not an incremental improvement. It establishes a new category of distributional RL algorithm — "Bellman sketch" methods — that is orthogonal to the categorical/quantile/expectile taxonomy. The framework doesn't just make existing sketch methods faster; it reveals that the entire distributional I/O layer was unnecessary, much as the realization that kernel methods could operate via the kernel trick without explicit feature maps changed how people thought about nonlinear classification.

Evidence anchor. Figure 11 (Appendix D.4) shows that Sketch-DP achieves better Cramér distance to ground-truth distributions than SFDP while running >100× faster per DP iteration — this empirically validates that eliminating imputation doesn't just save compute, it can actually improve distributional accuracy because the Bellman coefficients directly optimize the sketch-to-sketch mapping rather than routing through an imperfectly-inverted imputation. The theoretical bound in Proposition 4.4 shows that the approximation error shrinks as O(1/m), proving the approach is not merely a heuristic but a principled approximation scheme with controllable error.

---

Innovation 2: Approximate Bellman Closure via Least-Squares Regression — Escaping the Moment Straightjacket

The paper identifies and resolves a fundamental tension in sketch-based distributional RL that was previously treated as a binary property: a sketch is either Bellman-closed (exact dynamic programming possible) or not (requires imputation). Rowland et al. (2019, Theorem 4.3) characterized all Bellman-closed mean embedding sketches — they are limited to linear combinations of the first m moments, recovering Sobel's (1982) moment DP. This characterization implied that any sketch beyond low-order moments was necessarily incompatible with direct sketch-space Bellman backups, forcing the imputation round-trip.

What makes this distinctive. The paper shows that exact closure is unnecessarily strict. By defining Bellman coefficients as the solution to a least-squares regression problem (Equation 5) rather than requiring the equality φ(r + γg) = B_r φ(g) to hold pointwise, the framework achieves approximate Bellman closure for virtually any feature map. The key conceptual move is treating the closure condition as a regression target rather than a functional equation: instead of asking "does there exist B_r such that the equality holds for all g?", ask "what B_r minimizes the expected squared error over a weighting distribution μ?" This trivializes what was previously a restrictive classification theorem — any φ with linearly independent coordinates under μ yields well-defined B_r = C_r C^{-1}, and the quality of the approximation is quantified by the residual error.

Why the field was stuck. Prior work implicitly treated Bellman closure as a property of the sketch class: moments are closed, expectiles require imputation, quantiles require imputation. This created a sharp divide between algorithms that work in sketch space (only moments) and algorithms that must decode to distribution space (everything else). The moment-based approach was known to be numerically problematic at scale (the k-th moment grows as O(return^k), making TD learning with uniform step sizes impossible beyond m ≈ 5, as Appendix D.1 and Figure 7 empirically demonstrate), but there was no framework for getting the benefits of direct sketch-space updates (speed, linearity, theoretical tractability) for non-moment sketches.

What this unlocks. The regression-based approach transforms the design of distributional RL algorithms from a classification problem (is this sketch Bellman-closed or not?) into an approximation problem (how well can B_r approximate the bootstrap transformation under μ?). This has several important consequences:

• Rich feature design space. The translation family (Equation 8) — sigmoids, Gaussians, sinusoids, indicators — becomes usable for direct sketch-space DP/TD. These feature maps have much better numerical properties than moments: all coordinates are bounded (sigmoids in [0,1], Gaussians in (0,1]), localized (each feature responds to a specific region of the return range), and produce well-conditioned C matrices (unlike moment matrices that become exponentially ill-conditioned). Figure 7 concretely shows that a 50-dimensional sigmoid sketch vastly outperforms a 50-dimensional moment sketch in TD learning across all tested environments, with a wide basin of usable learning rates.

• Controllable approximation error. The regression residual directly bounds the per-step Bellman approximation error ε_B (Proposition 4.1), and by increasing the number of features m, one can drive this residual arbitrarily small for well-designed sketches (Proposition 4.4's O(1/m) bound for indicator features). This converts sketch-based distributional RL from an exact-or-nothing proposition to a consistency result: as m → ∞, the sketch values converge to the true mean embeddings. The framework thus provides the first asymptotic consistency guarantee for a non-moment sketch-based distributional RL method.

• Precomputation decouples feature design from runtime cost. The O(m^3) cost of computing B_r is paid once offline; after that, each backup costs O(m^2). This means the feature designer can freely explore high-dimensional feature maps without worrying about per-step computational overhead growing superlinearly, unlike imputation-based methods where imputation cost often scales poorly with m. In the deep RL experiments, using m = 401 features costs only 160,000 multiply-adds per backup — negligible relative to the convolutional forward pass — and the agent runs faster than QR-DQN and IQN (Table 1).

What kind of contribution this is. This is both a theoretical advance (replacing a binary classification with a continuous approximation framework with provable error bounds) and a practical enabler (making high-dimensional, numerically stable sketches usable for direct DP/TD). It is not a small tweak to Rowland et al. (2019) — it repurposes the regression problem at the heart of the Bellman coefficients from an implementational detail into the defining operation of a new algorithmic family. It also clarifies why moment-based methods failed in practice (not because moments are a bad representation per se, but because exact closure forced the use of globally-supported, widely-varying-scale features) and provides a constructive fix (use localized, bounded features with approximate closure).

Evidence anchor. Proposition 4.4 is the cleanest theoretical instantiation: for indicator features, the O(1/m) bound on asymptotic error follows directly from analyzing the three error sources (ε_B = Δ, ε_R = 2Δ, ε_E = 2Δ where Δ = (G_max - G_min)/m). Figure 4 empirically confirms that increasing m monotonically reduces both mean embedding error and Cramér distance to ground-truth distributions for a variety of feature maps (sigmoid, Gaussian, sinusoidal, indicator), with the excess Cramér distance shrinking toward zero — exactly what the theory predicts. The contrast with polynomial features (m = 50) in Figure 7, which perform catastrophically, underscores that the regression-based formulation is what makes non-moment sketches viable, not any particular choice of φ.

---

Innovation 3: Error Propagation Analysis as an Alternative to Operator Contractivity Proofs

Distributional RL convergence theory has historically been built around proving that the algorithmic operator — categorical projection composed with the distributional Bellman operator (Rowland et al., 2018), or quantile projection composed with the distributional Bellman operator (Dabney et al., 2018b; Rowland et al., 2023) — is a contraction in some probability metric (Cramér, Wasserstein-∞). This contractivity-then-fixed-point template is elegant when it works, but it fundamentally requires the algorithmic operator to be non-expansive in the chosen metric — a strong condition that fails for the sketch Bellman operator, whose singular values can exceed 1 (Appendix B.5, Figure 6).

What makes this distinctive. Instead of trying to prove that T^π_φ is a contraction (which it isn't), the paper develops an error propagation analysis that treats the sketch operator as an approximate simulator of the ideal operator Φ ◦ T^π. The analysis is structured around three sources of error — Bellman approximation (ε_B), reconstruction loss (ε_R), and embedding distortion (ε_E) — and shows that these errors compound sub-geometrically via the contraction of the true distributional operator T^π in Wasserstein/Cramér distance. The proof technique is visually illustrated in Figure 3: rather than analyzing T^π_φ in isolation, it traces a path through distribution space where contraction does hold, then measures how far the sketch operators deviate from this ideal path.

Why this is a conceptual advance, not just a different proof. The error propagation approach fundamentally changes what is required of the algorithmic operator. A contraction proof demands that T^π_φ map the space of sketch vectors to itself with a Lipschitz constant < 1 — a property that depends on the interaction between the Bellman coefficients, the transition dynamics, and the chosen norm, and that may simply not hold for many useful feature maps (as the singular value analysis suggests). The error propagation approach only requires that: (1) T^π is a contraction in some metric on distributions (true for many standard metrics); (2) the feature map provides bounded reconstruction and embedding error (a property of the sketch, independent of the MDP); (3) the Bellman coefficients have bounded regression error (a property of the offline fit, verifiable before running the algorithm). These conditions are modular and independently verifiable, unlike the monolithic contractivity condition.

What this enables for future work. This analysis template can be applied to any approximate operator derived from a contraction, not just sketch Bellman operators. If one designs a new distributional RL method by approximating some component of the Bellman backup (e.g., using a learned transition model, a compressed representation, or a reduced-rank Bellman coefficient), the same three-error-source decomposition applies, and the asymptotic bound follows mechanically. This makes the paper's theoretical contribution transferable — it provides a proof recipe, not just a proof for this specific algorithm.

The sharpest theoretical result: Proposition 4.4. The concrete bound for indicator features — lim sup ‖U_k − U^π‖_∞ ≤ (G_max − G_min)(3+2γ)/((1−γ)m) — is notable not just for its O(1/m) rate but for being fully explicit: every constant is expressed in terms of the MDP parameters (G_min, G_max, γ) and the number of features m. This is rare in distributional RL theory, where bounds typically involve unspecified constants or metric-dependent contraction factors. The explicit form lets a practitioner compute the required m for a desired accuracy before running the algorithm — m scales inversely with (1−γ), meaning longer horizons require more features, which makes intuitive sense since return distributions in low-γ environments are more concentrated and easier to approximate.

What kind of contribution this is. This is a theoretical innovation — a new proof technique that solves a problem (non-contractive operators) that the standard approach couldn't handle. It's also a diagnostic tool: the three-error decomposition (ε_B, ε_R, ε_E) provides a vocabulary for discussing why a sketch-based method might fail, distinguishing between poor feature design (large ε_R, ε_E) and poor Bellman coefficient fit (large ε_B). The paper's empirical sweeps over m and slope s (Figure 4) can be reinterpreted through this lens: increasing m reduces all three errors by providing finer coverage; intermediate slopes minimize ε_B by matching feature smoothness to the bootstrap transformation, while extremes increase it.

Limitation to note. The theory covers only the DP case; convergence of Sketch-TD is explicitly left as future work (Section 7). The DP analysis provides a template (treat the TD update as a stochastic approximation to the DP operator and analyze the resulting dynamical system), but the stochastic case introduces additional technical challenges — Martingale noise, off-policy sampling, function approximation — that the paper does not address. This is a significant gap between the theoretical guarantees and the practical algorithms (Sketch-TD, Sketch-DQN) that motivated the work.

---

Innovation 4: Verifier-Independent Distributional Learning as a Design Axis

Distributional RL methods are typically categorized by how they approximate distributions: categorical methods discretize the CDF support, quantile methods discretize the probability levels, expectile methods learn specific functionals. The Bellman sketch framework introduces a new axis of variation orthogonal to this taxonomy: the choice of feature map φ determines which statistical properties of the return distribution are preserved by the representation. This is a fundamentally different design philosophy from "choose a distribution family and minimize a statistical distance."

What the field did before. In categorical DQN (Bellemare et al., 2017), the designer chooses the support grid {z_1, ..., z_m} — this determines the resolution and range of the distribution representation, but the algorithm always learns a full categorical distribution (all bin probabilities). In QR-DQN (Dabney et al., 2018b), the designer chooses the number of quantiles m, but the algorithm always learns uniform-weight quantiles — it cannot selectively focus on tails, modes, or specific distributional features. IQN (Dabney et al., 2018a) makes the quantile levels themselves learnable, but still represents the full quantile function. These methods differ in parametrization but share a common goal: approximate the entire return distribution as accurately as possible given the representational budget.

What makes this distinctive. The Bellman sketch framework decouples what you learn about the distribution from how you parametrize it. Choosing a Gaussian feature map with anchors concentrated in the left tail means the sketch will capture fine details of downside risk while coarsely representing the right tail. Choosing sinusoids means the sketch captures periodic structure in the return distribution (relevant if rewards have cyclical patterns). Choosing indicator features with a uniform grid recovers something like categorical distributional RL but with sketch-space updates. The feature map φ is a declarative specification of the distributional features of interest — it's a prior over what matters about the return distribution, rather than a neutral representation budget.

Why this matters. In risk-sensitive RL, an agent may care only about the 5th percentile of returns (Value at Risk) or the expected shortfall below a threshold. A categorical or quantile method with m = 200 must allocate representation capacity uniformly across the entire distribution to get the tail right — most of the capacity is wasted on the bulk of the distribution. A sketch with sigmoid features concentrated in the left tail can achieve better tail accuracy with far fewer features because it focuses capacity where it matters. This targeted representation is not achievable in the standard categorical/quantile frameworks without modifying the loss function or adding auxiliary objectives. The value readout coefficients β further support this philosophy: they extract the expected value linearly from whatever distributional features φ captures, but one could equally extract other functionals (variance, quantiles, expectiles) by solving different readout regression problems — all reusing the same learned sketch values.

The negative result that clarifies the contribution. Appendix K's ReST^EM experiment is brief but important: attempting to optimize the sketch-based revision model with on-policy RL training degraded performance. While this is in the context of a different paper architecture, it illustrates a general point that the paper's framework makes visible: the choice of φ imposes an inductive bias about what distributional structure matters, and training procedures that ignore this bias (by, e.g., optimizing the sketch to predict something φ wasn't designed to capture) can backfire. The framework thus serves as a diagnostic tool for understanding when and why distributional RL methods succeed or fail — a role that goes beyond proposing yet another distribution representation.

What kind of contribution this is. This is a conceptual reframing of the distributional RL design space, adding a new axis (feature map choice as a target specification) that is independent of existing axes (parametrization family, loss function, optimization algorithm). It is not a performance claim — the Atari results in Figure 5 show Sketch-DQN performing comparably to existing methods, not dominating them — but rather an expansion of what kinds of algorithms can be designed and what design questions practitioners should ask. The paper's systematic sweeps over base features (sigmoid, Gaussian, parabolic, tanh, sinusoidal, indicator), feature counts m, and slope s (Figure 4, Figure 9) are best understood as an empirical exploration of this new design space, demonstrating that the framework is not a single algorithm but a generator of algorithms parameterized by φ.

Evidence anchor. Figures 9a–d show that the choice of feature map and slope produces qualitatively different tradeoffs between mean embedding accuracy and Cramér distance, and that no single configuration dominates across all environments — exactly what one would expect if φ encodes a prior over distributional features. The deep RL experiments (Figure 13) further show that sigmoid features outperform Gaussian features for the Atari suite under the chosen hyperparameters, confirming that feature map selection matters for downstream task performance, not just for reconstruction fidelity. The ablation in Figure 10 shows that the anchor range (relative to the return range) affects accuracy in systematic ways — too narrow loses probability mass, too wide loses resolution — providing actionable guidance for practitioners using the framework.

5. Experimental Analysis

Evaluation Methodology

• Dataset. All tabular experiments use custom Markov reward processes (MRPs) described in Appendix C.1: Random chain (10 states, stochastic transitions, deterministic reward of +1 at state 10), Directed chain or DC (5 states, deterministic transitions, deterministic reward of +1 at state 5), DC+Gaussian R (same as DC but with Gaussian reward 𝒩(1, 1) at state 5), Tree, Loopy tree, Cycle, Rowland '23 (Example 6.5 of Rowland et al., 2023), S&B '18 (Example 6.4 of Sutton & Barto, 2018), and Loopy fork (Figure 2A). All environments use discount factor γ = 0.9. For deep RL, the benchmark is the Atari 57 suite (Bellemare et al., 2013) with standard 200M frame training, evaluated using human-normalized scores (Mnih et al., 2015).

• Base model(s). Tabular experiments assume full knowledge of the MDP transition and reward structure (for DP) or sample transitions from the true MDP (for TD). Deep RL experiments use a neural network architecture matching QR-DQN (Dabney et al., 2018b): a convolutional torso feeding into a fully-connected layer that outputs m sketch values per action, with a sigmoid nonlinearity on the final output to bound sketch predictions to [0, 1]. The architecture is not a pretrained model — it is trained from scratch for 200M frames using the standard DQN training procedure (Mnih et al., 2015).

• Metrics. Three categories of metrics are used:

  • Mean embedding squared error: the squared L₂ distance ‖U_est(x) − U^π(x)‖₂² between the learned sketch values and ground-truth mean embeddings, where ground truth is estimated from 10⁵ Monte Carlo returns per state with horizon 110 (deterministic rewards) or 200 (Gaussian rewards) after first visit, producing truncation error bounded at 10⁻⁴.
  • Cramér distance: the ℓ₂ distance between the CDF of the distribution imputed from learned sketch values and the CDF of the ground-truth return distribution, reported as max_{x∈𝒳} ℓ₂(η̂(x), η^π(x)), where ground-truth distributions are projected onto the same (jittered) support as the imputed distributions to isolate the sketch's representational error from discretization error. Reported as the average over 100 independent jitters of the support (Appendix C.1).
  • Excess Cramér distance: the Cramér distance minus the Cramér distance between the categorical projection of η^π onto the support grid and η^π itself — this measures how much additional error the sketch introduces beyond what is unavoidable from the support discretization alone.
  • Deep RL performance: median and mean human-normalized returns across all 57 Atari games, computed from 3 random seeds per game.

• Baselines.

  • Categorical DP (CDRL): standard categorical distributional RL with fixed particle locations and learnable probabilities (Rowland et al., 2018; Bellemare et al., 2017), using the same number of bins as the sketch has features for fair comparison.
  • Dirac delta at mean return: a naive baseline that places all probability mass at V^π(x), representing a non-distributional method. Included to verify that sketches capture meaningful distributional information beyond the mean.
  • SFDP: Statistical functional dynamic programming (Rowland et al., 2019) with expectile sketches, using imputation strategies to round-trip through distribution space, compared against Sketch-DP for computational efficiency and accuracy (Appendix D.4, Figure 11).
  • DQN: standard non-distributional deep Q-network (Mnih et al., 2015).
  • C51: categorical distributional DQN with 51 atoms (Bellemare et al., 2017).
  • QR-DQN: quantile regression DQN with 201 quantiles (Dabney et al., 2018b), tuned to best-performing quantile count.
  • IQN: implicit quantile networks with 51 quantiles (Dabney et al., 2018a), tuned to best-performing quantile count.

• Generation budget / compute accounting. For tabular DP experiments, "compute" is measured in DP iterations — all methods are run for 200 iterations (well past convergence). For the wallclock comparison in Figure 11, per-iteration time and one-time operator setup time are measured directly. For deep RL, compute is measured in training frames (200M for all agents), with frame processing rates reported in Table 1 to control for per-step computational cost differences. The Bellman coefficient precomputation cost (O(m³) matrix inversion) is a one-time offline cost and is not amortized over the per-iteration or per-frame measurements, though Figure 11 does report setup time separately.

• Cross-validation / statistical protocol. No cross-validation is used in the tabular experiments — ground-truth return distributions are estimated from independent Monte Carlo rollouts. For the deep RL experiments, each agent is run with 3 random seeds per game, and performance is reported as the mean/median across all games and seeds. The paper does not report confidence intervals or statistical significance tests. For the Sketch-DQN hyperparameter sweep (Appendix C.2), the slope s was tuned from {1, ..., 12} and the feature count m from {101, 201, 401}, but the selection criterion was the worst-case regression error (reject if > 0.01) rather than a held-out validation set.

Main Quantitative Results

Tabular DP: Feature Count and Slope Sweeps

Headline result. Across three MRP environments (Random chain, Directed chain, DC+Gaussian R) and seven feature map types (sigmoid, Gaussian, parabolic, tanh, sinusoidal, indicator, and CDRL baseline), increasing the number of features m monotonically reduces both mean embedding squared error and Cramér distance to ground truth, with excess Cramér distance decreasing toward zero — empirically confirming Proposition 4.4's theoretical prediction that arbitrary accuracy is achievable as m → ∞ (Figure 4, first row).

Feature count sweep (m = 10, 50, 90). At m = 10, mean embedding squared error is on the order of 10⁻⁴ to 10⁰ depending on the environment and feature type (Figure 4, left column). By m = 90, all feature types except the CDRL baseline reduce this error by 1–4 orders of magnitude. For the Gaussian-based sketch on the Random chain, the error drops from approximately 10⁻² at m = 10 to approximately 10⁻⁶ at m = 90. The Cramér distance follows a similar pattern: all sketch types start at Cramér distances of 0.1–0.5 at m = 10 and decrease to 0.05–0.2 at m = 90, approaching the CDRL reference line. The excess Cramér distance — which isolates the sketch-specific error beyond the categorical discretization — drops to 10⁻⁷–10⁻² at m = 90, indicating that the sketch representation becomes nearly as faithful to the return distribution as the categorical representation at equal support resolution.

Comparison against CDRL. CDRL (shown as a separate curve/bar in each plot) generally achieves the lowest or near-lowest Cramér distances, consistent with its direct optimization of distributional fit. However, several sketch configurations match or approach CDRL performance at high m. For example, on the Directed chain at m = 90, sigmoidal and Gaussian sketches both reach Cramér distances within ~0.02 of CDRL, while operating entirely in sketch space without representing distributions. On DC+Gaussian R, CDRL achieves a Cramér distance of approximately 0.2 at m = 90, while Gaussian sketches reach approximately 0.25 — a gap that may be acceptable given the sketch's faster per-iteration cost and flexibility.

Dirac baseline. The Dirac at mean return is included as a dotted reference line in the Cramér distance plots. All distributional methods (including Sketch-DP with m ≥ 50) significantly outperform the Dirac baseline on all environments, confirming that the sketch captures meaningful distributional structure beyond the mean. For instance, on the Random chain, the Dirac Cramér distance is approximately 0.3, while sigmoid sketches at m = 50 achieve approximately 0.15.

Slope sweep (scaling factor 0.001 to 10.0 relative to default slope). The slope parameter s — controlling the sharpness of the translation family features — exhibits a U-shaped or valley-shaped dependence for most feature types (Figure 4, second row). Mean embedding error is minimized at large slopes (sharp features) for sigmoid, Gaussian, and parabolic features, often by orders of magnitude relative to small slopes. This is expected: sharper features can localize more precisely, reducing regression error ε_B. However, Cramér distance shows a different pattern: it is minimized at intermediate slopes (scaling factor 0.1–1.0), and rises at both extremes. When features are too sharp (large s), they become near-indicator functions that only respond in tiny intervals — there exist regions of the return range where no feature varies meaningfully, producing "dead zones" where the sketch is uninformative about the distribution. When features are too smooth (small s), they are nearly constant over the entire return range, providing little distributional resolution. The optimal slope balances these two failure modes. For sinusoidal features, the trends are less pronounced, with relatively flat Cramér distance across slope values.

Feature type comparisons. Sigmoid and Gaussian features achieve the best Cramér distances at high feature counts, with sigmoid performing slightly better in most environments. Parabolic features (κ(x) = 1 − x² for |x| ≤ 1) perform comparably but with more variance across environments. Hyperbolic tangent features perform slightly worse. Sinusoidal features, despite being non-local (global oscillations), achieve competitive performance on some environments but underperform on others — notably, the Cramér distance for sinusoidal features on DC+Gaussian R does not decrease as reliably with m as for localized features. This aligns with the intuition that localized features match the local structure of return distribution CDFs better than global features.

Extended environments (Figure 9). The additional environments (Tree, Loopy tree, Cycle, Rowland '23, S&B '18, Loopy fork) with both deterministic and Gaussian rewards broadly replicate the findings. For deterministic-reward environments (Figure 9a–b), the feature count and slope sweeps show qualitatively identical trends to the main three environments. For Gaussian-reward environments (Figure 9c–d), the mean embedding error is generally higher (stochastic returns produce broader distributions that stress the sketch representation), but the monotonic improvement with m and the U-shaped slope dependence persist. CDRL's advantage over sketches is slightly larger in Gaussian-reward environments, suggesting that explicit distributional representations may have an edge when return distributions are smooth and unbounded, but the sketch methods still substantially outperform the Dirac baseline.

Anchor range ratio sweep (Figure 10). Sweeping the ratio of the anchor range width to the regression grid μ's support width (from 0 to 3) reveals a sharp asymmetric sensitivity. When the anchor range ratio drops below 1 (anchors narrower than the regression grid → narrower than the true return range), Cramér distances increase sharply because the support points for imputation may miss significant probability mass outside the anchor range. When the ratio increases above 1 (anchors wider than the return range), the Cramér distance increases more gently because anchor points become sparser, reducing imputation resolution but not missing probability mass. The optimal ratio is slightly above 1 for most environments and feature types, consistent with the heuristic recommendation in Appendix B.6. This validates the practical guidance: "choose the anchor range to be slightly wider than the return range."

Tabular DP: Sketch-DP vs. SFDP Comparison

Headline result. Sketch-DP achieves better Cramér distances than SFDP (expectile-based, with imputation) while running >100× faster per DP iteration across all tested environments and feature/expectile counts m ∈ {25, 50, 75} (Figure 11, rows 1–3).

Accuracy comparison (Figure 11, rows 1–2). For m = 25, Sketch-DP and SFDP achieve comparable Cramér distances (within ~0.02) on most environments. As m increases to 75, Sketch-DP Cramér distances continue to decrease (reaching 0.05–0.15 on most environments), while SFDP Cramér distances either plateau or increase on several environments (Loopy tree, Rowland '23). The excess Cramér distance shows the same pattern more clearly: Sketch-DP excess error drops to 10⁻⁴–10⁻² at m = 75, while SFDP excess error sometimes increases with m (e.g., on Rowland '23, SFDP excess error at m = 75 is larger than at m = 25). This suggests that the imputation strategy in SFDP becomes less accurate as the number of expectiles grows — the optimization problem for reconstructing a distribution from more expectile values may become more ill-conditioned or harder to solve to high precision, eroding the benefit of additional features.

Computational efficiency (Figure 11, rows 3–4). The per-iteration wallclock time for Sketch-DP is below 10⁻² seconds for all m values and environments. For SFDP, per-iteration time ranges from ~0.2 seconds (m = 25) to ~1–4 seconds (m = 75), representing a speedup of 100–400× for Sketch-DP. The one-time operator setup time (computing Bellman coefficients for Sketch-DP, setting up the imputation optimization for SFDP) is 0.1–4 seconds for Sketch-DP (depending on m) versus negligible for SFDP, but this cost is amortized over the full DP run. Since SFDP requires ~200 iterations, its total runtime is 40–800 seconds, while Sketch-DP's total is 0.1–4 seconds setup plus <2 seconds for iteration — a practical difference of 20–400× in total runtime.

Tabular TD: Polynomial vs. Sigmoid Features

Headline result. Sigmoid-based Sketch-TD with m = 50 features finds a wide basin of usable learning rates (roughly 10⁻⁴ to 10⁻²) and achieves Cramér distances of 0.05–0.3 across environments, while moment-based sketches with m = 50 polynomial features fail to find any learning rate achieving Cramér distance below ~1.0 — performing worse than the Dirac baseline on most environments (Figure 7).

Learning rate sensitivity (Figure 7). For sigmoid features (m = 50), the Cramér distance curve is U-shaped in learning rate, with a wide flat minimum spanning ~2 orders of magnitude (e.g., 10⁻⁴ to 10⁻² on Random chain, 10⁻³ to 10⁻² on Directed chain). For m = 5 polynomial features, there is also a usable basin, but the minimum Cramér distance is substantially worse than sigmoid (e.g., ~0.3 vs. ~0.1 on Random chain). For m = 50 polynomial features, the Cramér distance is above 0.8 at all learning rates tested (10⁻⁶ to 1.0) on most environments — far worse than either m = 5 polynomials or m = 50 sigmoids, and often worse than the Dirac baseline. This empirically confirms the paper's claim that moments "are naturally of widely differing magnitudes" making "a single learning rate impossible to tune" (Section 3.1, Appendix D.1).

Why this matters for the framework's motivation. The failure of m = 50 polynomial sketches is the empirical justification for why Bellman closure alone is insufficient — exact DP is possible for moments, but the numerical instability prevents TD learning at scale. The success of m = 50 sigmoid sketches (which are not Bellman-closed but use approximate regression-based Bellman coefficients) validates the paper's central thesis: approximate closure via least-squares is not just a theoretical convenience but a practical necessity for high-dimensional sketch learning.

Deep RL: Atari 57 Benchmark

Headline result. Sketch-DQN (sigmoid features, m = 401) achieves median human-normalized return of ~1.7 and mean ~12.5 at 200M frames, outperforming C51 and QR-DQN, approaching IQN's performance, while running faster than both QR-DQN and IQN (Figure 5, Table 1).

Learning curves (Figure 5, left: median; right: mean). At 200M frames (the standard Atari benchmark horizon), the median human-normalized scores are:

• DQN: ~0.8
• C51 (51 atoms): ~1.2
• QR-DQN (201 quantiles): ~1.4
• IQN (51 quantiles): ~1.9
• Sketch-DQN (401 sigmoid features): ~1.7

The mean human-normalized scores are:

• DQN: ~4.5
• C51: ~8.0
• QR-DQN: ~10.5
• IQN: ~14.0
• Sketch-DQN: ~12.5

Sketch-DQN's mean score falls between QR-DQN and IQN; its median score is above QR-DQN and C51 but below IQN. The gap between mean and median is driven by a few high-scoring games where Sketch-DQN excels — the per-game advantage plot (Figure 12) shows consistent advantages over DQN, C51, and QR-DQN on games like CRAZY CLIMBER, SPACE INVADERS, RIVER RAID, ROAD RUNNERS, and VIDEO PINBALL, while modestly underperforming on ASSAULT, ASTERIX, DOUBLE DUNK, KRULL, PHOENIX, and STAR GUNNER.

Comparison to IQN. The paper notes that IQN "uses a more complex prediction network to make non-parametric predictions of the quantile function" (Section 5.1). IQN requires a separate forward pass through the MLP layers for each quantile level, making it slower per frame (~1120 fps at 64 quantiles, dropping to ~698 fps at 201 quantiles; Table 1) and architecturally more complex. Sketch-DQN produces all 401 sketch values from a single final layer, achieving 1326 fps. The performance gap (median ~1.7 vs. ~1.9; mean ~12.5 vs. ~14.0) must be weighed against this computational simplicity.

Frame processing rates (Table 1). On a single V100 GPU, average frames per second:

• Sketch-DQN (m = 401): 1326 ± 107
• C51 (51 atoms): 1309 ± 146
• QR-DQN (201 quantiles): 1258 ± 107
• QR-DQN (64 quantiles): 1286 ± 103 (from trend in the table)
• IQN (64 quantiles, default): 1120 ± 90
• IQN (201 quantiles): 698 ± 41
• IQN (401 quantiles, extrapolated): ~400 ± 16

Sketch-DQN is the fastest distributional agent, with C51 and QR-DQN slightly slower. IQN scales poorly with quantile count — at m = 64 (default), it is already slower; at m = 201, its frame rate drops by ~40% relative to QR-DQN at the same count. The sketch architecture's single-forward-pass design keeps it efficient even at m = 401.

Feature parameter sensitivity (Figure 13). Ablating over feature type (sigmoid vs. Gaussian), feature count m (101 vs. 201), and slope s reveals:

• Sigmoid features consistently outperform Gaussian features at equal m and comparable s. At m = 201, sigmoid (s = 5) achieves mean ~12.5 while Gaussian (s = 1.67) achieves ~10–11.
• Increasing m from 101 to 201 improves performance for sigmoid features (mean from ~11.5 to ~12.5) but has minimal impact for Gaussian features, suggesting Gaussian features at the tested slopes are saturating in representational capacity earlier.
• The slope tuning is important within a feature family: for sigmoid at m = 201, s = 5 outperforms s = 4 and s = 6; for Gaussian at m = 201, the effect is less pronounced.

Value readout β. The β coefficients enable Q-learning by linearly mapping sketch values to expected returns. The paper does not ablate this choice (e.g., comparing to learning the value directly from a separate network head), which would have clarified whether the two-stage approach (learn sketch, read out value) is beneficial beyond enabling Q-learning to function.

Ablation Studies and Robustness Checks

• Feature count m (Figure 4, Figure 9): Increasing m monotonically improves mean embedding accuracy and Cramér distance for all tested feature maps across all environments, with diminishing returns beyond m ≈ 50–90. Excess Cramér distance decreases with m, confirming the sketch convergences to the same distributional fidelity as the categorical representation at equal resolution. This is the primary empirical validation of Proposition 4.4's O(1/m) theoretical bound.

• Slope s (Figure 4, Figure 9): The U-shaped dependence of Cramér distance on slope (minimized at intermediate values, degraded at extremes) is robust across feature types and environments, though the optimal range shifts by feature. For sigmoid and Gaussian features, the default slope heuristic (50% overlap, 10 features covering the return range; Appendix B.6) falls within the optimal region, validating the design guidance empirically. Sinusoidal features show flatter dependence, consistent with their global nature.

• Anchor range ratio (Figure 10): Anchors slightly wider than the regression grid (ratio 1.0–1.5) are optimal. Ratios below 0.8 cause sharp Cramér distance degradation because imputation support misses probability mass; ratios above 2.0 cause gentle degradation from reduced resolution. This is robust across environments and feature types, providing a concrete hyperparameter guideline.

• Base feature κ (Figures 4, 9): No single κ dominates across all metrics and environments. Sigmoid and Gaussian features reliably perform well; hyperbolic tangent slightly worse; parabolic features more variable; sinusoidal features competitive on some environments but less reliable. The choice of κ is impactful and the paper's framework treats it as a design degree of freedom, not a parameter to optimize automatically.

• Imputation vs. direct sketch evaluation (Figure 4 vs. Figure 8): The paper evaluates sketches both by direct mean embedding error (which measures how accurately the sketch values themselves are estimated) and by Cramér distance after imputation (which measures how much distributional information is retained). The two metrics sometimes diverge — e.g., Gaussian features at large slope achieve very low mean embedding error but higher Cramér distance — illustrating that mean embedding accuracy alone does not guarantee faithful distributional recovery, and that feature choice involves trading off these two aspects.

• SFDP vs. Sketch-DP (Figure 11): The wallclock speedup (100–400× per iteration) and accuracy improvement (lower Cramér distances at m = 75) over SFDP validates the paper's central claim that eliminating imputation is both faster and potentially more accurate, since the Bellman coefficients directly optimize the sketch-to-sketch mapping rather than routing through an imperfectly inverted imputation. The fact that SFDP's accuracy sometimes degrades with more expectiles (m = 75 worse than m = 50 on some environments) is an unanticipated negative result that underscores the fragility of imputation-based approaches at scale.

• Polynomial vs. bounded features for TD (Figure 7): The catastrophic failure of m = 50 polynomial features across all tested learning rates, contrasted with the robust performance of m = 50 sigmoid features, is a clean ablation demonstrating that the issue with moment-based sketches is numerical (divergent feature scales) rather than representational (insufficient capacity). This justifies the paper's move from exact Bellman closure (moments) to approximate closure (general φ).

• Sigmoid vs. Gaussian features in deep RL (Figure 13): In the Atari setting, sigmoid features outperform Gaussian features by 1–2 mean normalized return points at comparable configurations. Within sigmoid, increasing m from 101 to 201 yields clear improvement; within Gaussian, increasing m has negligible effect. The paper does not ablate m beyond 201, leaving open whether further increases would close the gap.

• Output nonlinearity in Sketch-DQN (Appendix C.2, described textually): Adding a sigmoid or tanh nonlinearity to the final network layer to bound sketch predictions to the known output range of φ was "found to be very crucial for a good performance" — though explicit ablation curves are not shown.

• Regularization in Bellman coefficient computation (Appendix C.2): L₂ regularization with weight 10⁻⁹ was tuned from {10⁻¹⁵, 10⁻¹², 10⁻⁹, 10⁻⁶, 10⁻³} to avoid numerical issues from near-singular C. Larger regularization was rejected because it increases ε_B (worse Bellman approximation). The chosen small value indicates that C is well-conditioned for the tested feature configurations.

• Rejection of high-slope configurations (Appendix C.2): In the deep RL tuning, configurations where the worst-case regression error max_r max_g ‖φ(r + γg) − B_r φ(g)‖ exceeded 0.01 were rejected upfront, without running full agent training. This is a practical application of Proposition 4.1's insight that the per-step Bellman approximation error is bounded by this quantity — the paper uses it as a screening criterion to avoid training agents with predictably poor Bellman coefficients.

Critical Assessment

The experimental evaluation provides substantial evidence for the paper's core methodological claims but has specific limitations that constrain the strength and generality of the conclusions.

On the claim that Sketch-DP achieves arbitrary accuracy as m → ∞ (Proposition 4.4). The experiments support this qualitatively: mean embedding error and Cramér distance monotonically decrease with m for all feature types and environments (Figures 4, 9). However, the paper does not empirically verify the O(1/m) rate — the sweeps test only three m values (10, 50, 90), which is insufficient to distinguish O(1/m) from O(1/√m) or O(1/log m) empirically. The theoretical rate is validated only analytically, not experimentally. Moreover, the imputation step used to compute Cramér distances introduces its own approximation error (solving a quadratic program on a finite support), which could confound the measured scaling. The excess Cramér distance metric partially controls for this by subtracting the categorical projection error, but the imputation from sketch values may still be less accurate than the categorical projection from ground-truth distributions, especially at small m. A stronger empirical validation would measure the distance between sketch values directly (‖U_k − U^π‖, which requires no imputation) and confirm the predicted linear dependence on 1/m.

On the claim that the framework is general across feature maps. The experiments test an impressive variety of feature maps — sigmoid, Gaussian, parabolic, tanh, sinusoidal, indicator, and polynomial — across many MRPs. The consistent qualitative trends (m improves accuracy, intermediate slope is optimal) support the claim that the framework is robust to φ choice. However, the quantitative differences between feature maps are substantial and not fully explained. Why do sigmoid features achieve Cramér distance ~0.05 on some environments at m = 90 while sinusoids achieve ~0.15? The paper attributes this to localization (local features match CDF structure better) but does not test hybrid designs (e.g., multi-scale features, learned anchor placements, features with varying slopes) that could combine the benefits. The framework's claim to generality is supported in the sense that "many φ work," but not in the sense that "φ choice is well-understood" — the practitioner is still left with a hyperparameter tuning problem for each new domain.

On the claim that Sketch-DP is faster than SFDP. Figure 11 provides clean evidence: 100–400× per-iteration speedup across environments and m values. This is the most unambiguous experimental result in the paper. However, the comparison is against a specific SFDP implementation (expectile sketches, SciPy's MINIMIZE for imputation), and the SFDP runtime could potentially be improved with custom optimizers, warm-starting from previous imputations, or approximate imputation strategies. The paper's claim that eliminating imputation is faster is well-supported; the specific 100× factor is implementation-dependent.

On the claim that the sketch approach enables tractable TD learning where moments fail. Figure 7 is a clean, convincing ablation: m = 50 sigmoid features find good learning rates across environments; m = 50 polynomial features fail at all learning rates. This directly validates the numerical instability argument that motivates the approximate-closure approach. However, the comparison is slightly unfair: the polynomial features use unnormalized monomials 1, g, g², ..., g⁴⁹, which are known to be ill-conditioned. Standard practice would normalize or orthogonalize these features (e.g., use Legendre or Chebyshev polynomials, or apply standardization), which could potentially rescue moment-based TD learning. The paper does not test whether orthogonalized polynomial features with Bellman coefficients computed via regression (as in the general sketch framework) would perform better — if they did, the distinction between "moment methods" and "general sketches" would blur, with the key innovation being the regression-based Bellman coefficients rather than the feature type per se.

On the claim that Sketch-DQN approaches IQN performance while being faster. The Atari results (Figure 5) support this claim with qualifications. Sketch-DQN (median ~1.7, mean ~12.5) does outperform C51 and QR-DQN and approach IQN (median ~1.9, mean ~14.0). The runtime advantage (Table 1: 1326 fps vs. IQN's 1120 fps at default 64 quantiles, or 698 fps at 201 quantiles) is clear. However, IQN with 64 quantiles is IQN's default configuration, not necessarily its Pareto-optimal configuration in the accuracy–speed tradeoff. IQN with 64 quantiles at 1120 fps achieves better performance than Sketch-DQN with 401 features at 1326 fps — the 15% speedup comes with slightly lower performance. A fairer comparison would plot a Pareto frontier: Sketch-DQN at various m (101, 201, 401) vs. IQN at various quantile counts (8, 16, 32, 64, 128), measuring both median return and fps. The paper's point that the sketch architecture is simpler (single forward pass) is valid, but the performance claim is relative to IQN's default configuration, not to IQN's best configuration at equal compute.

On the broader claim that the framework is "principled" and "rigorous." The DP convergence theory (Section 4) is rigorous and novel. The TD convergence theory is explicitly absent — the paper states "convergence analysis for general sketches is an immediate future work" (Section 7). This means the theoretical guarantees cover only the tabular DP case, while the motivating applications (deep RL, TD learning) operate in the theoretically unanalyzed regime. This is a significant gap between theory and practice that the paper acknowledges honestly but does not bridge. The deep RL results are thus empirical demonstrations that the approach can work in practice, not instantiations of a theoretically guaranteed method.

Missing experiments that would have strengthened the paper:

1. Direct comparison of Bellman coefficient regression error ε_B vs. measured DP error. Proposition 4.1 bounds the per-step error by the worst-case regression residual. The paper could compute this residual for each feature configuration (since φ, B_r, and μ are all known) and plot it against the measured asymptotic error ‖U_∞ − U^π‖ to empirically verify the error propagation bound and identify which of the three error sources (ε_B, ε_R, ε_E) dominates in practice.

2. Sketch-DP with very high m (200, 500, 1000). The sweeps stop at m = 90. Given the O(m³) precomputation cost, testing m = 500 on small MRPs would be feasible and would provide stronger evidence for the asymptotic scaling. The theoretical O(1/m) bound predicts continued improvement; confirming or refuting this at higher m would be valuable, especially since the excess Cramér distance in Figure 4 appears to be approaching zero (or a floor) at m = 90.

3. Sketch-TD convergence curves with varying learning rates for more feature types. Figure 7 shows polynomial vs. sigmoid at m = 5 and m = 50. Including Gaussian, sinusoidal, and indicator features at m = 50 in the same learning rate sweep would test whether the robust learning rate basin is specific to sigmoid features or general to bounded, well-scaled features. The paper's claim that the framework "generalizes" to arbitrary φ would be strengthened by showing that multiple feature types produce stable TD learning at m = 50.

4. Ablation of the constant feature. The paper states that appending a constant feature was "very crucial for a good performance" but never shows what happens without it — an ablation that would take one additional curve per plot. The constant feature converts the sketch operator from linear to affine, which is a fundamental structural choice. Quantifying its impact (does performance collapse entirely? degrade moderately?) would clarify the importance of the affine structure.

5. Comparison against learned-feature baselines. The sketch framework treats φ as fixed and pre-designed. Modern deep RL often learns representations end-to-end. A comparison against a method that learns φ (or learns the Bellman coefficients directly from data, without the precomputation-on-μ step) would position the paper's approach relative to fully learned alternatives. The paper acknowledges the possibility of learned features (Remark 3.2, Appendix B.4) but does not empirically explore it.

6. Statistical error bars on Cramér distances. The imputation step uses 100 independent support jitters, and the paper reports the average. Reporting standard deviations would indicate whether the observed differences between feature types at high m are statistically meaningful or within the noise of the imputation procedure.

On the conditional nature of the Atari results. The Sketch-DQN performance depends on careful tuning of feature type, m, slope, and regularization — the paper found that Gaussian features underperform sigmoid, that m = 401 is needed for best performance, and that configurations with ε_B > 0.01 are rejected. This is a substantially larger hyperparameter space than C51 (which tunes only the number and range of atoms) or QR-DQN (which tunes only the number of quantiles). The framework's flexibility comes at the cost of a larger tuning burden, and the paper does not provide a principled method for selecting φ beyond the heuristics in Appendix B.6. The Atari results demonstrate feasibility (a well-tuned Sketch-DQN can be competitive) but not robustness (an arbitrary φ with default settings may perform poorly).

Overall assessment. The experimental evaluation effectively validates the paper's primary methodological contributions: (1) the sketch Bellman operator enables accurate, efficient DP without imputation; (2) the regression-based Bellman coefficients enable non-moment sketches to work at scale where exact moment methods fail; (3) the framework scales to deep RL with competitive performance. The experiments are thorough within their scope (many feature types, many MRPs, standard Atari benchmark) but leave open questions about asymptotic scaling rates, TD convergence theory, feature selection principles, and the tradeoff between framework flexibility and tuning complexity. The largest gap is between the DP theory (rigorous convergence) and the deep RL practice (empirical only), which the paper acknowledges as future work.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted For and Dominates the Headline Efficiency Gains

The assumption or constraint. The entire compute-optimal framework depends on estimating each prompt's difficulty before deciding how to allocate the inference budget. The paper's method for doing so — generating 2048 samples per question and computing the pass@1 rate (oracle) or averaging the PRM's final-answer scores (predicted) — is extraordinarily expensive. The authors acknowledge this explicitly 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"

The consequence. The headline $4\times$ efficiency gain over best-of-N (Figures 4 and 8) is computed after difficulty is known, without amortizing the cost of learning it. Generating 2048 samples per question to estimate difficulty consumes more compute than the largest test-time budgets studied (256–512 generations). If difficulty estimation cost were included, the total cost would be difficulty estimation + strategy execution, and the former could dominate. A practitioner implementing this approach would find that the wall-clock time and FLOP count per question are dominated by the 2048-sample pre-evaluation, not by the optimized strategy, making the $4\times$ figure an upper bound on achievable efficiency rather than a realized deployment gain. On a 500-question test set, the difficulty estimation alone costs 500 × 2048 = 1,024,000 generations — equivalent to running best-of-2048 on every question, which is far more expensive than any strategy the compute-optimal policy deploys.

What evidence exists in the paper. Section 3.2 describes the 2048-sample procedure. The paper does not report any experiment where difficulty estimation cost is included in the compute budget, nor does it compare total cost (estimation + execution) against a uniform best-of-N baseline with the same total budget. The theoretical framework in Section 3.1 defines the optimization over $\theta$ given a budget $N$, but $N$ is defined as the generation budget for solving the problem, not including any pre-evaluation. The paper acknowledges this gap but provides no empirical characterization of the tradeoff.

Mitigation status. The paper explicitly flags this as a key avenue for future work (Section 3.2: "We flag that reducing the cost of this step is an important avenue for future work, and note that it can naturally be cast as an exploration-exploitation trade-off"). Section 8 further suggests "pretraining or finetuning models to directly predict difficulty of a question" or "training a difficulty predictor offline." However, no such model is developed or evaluated in the paper. The mitigation is entirely aspirational — the current method provides no practical solution for deploying compute-optimal scaling without prohibitive pre-evaluation cost.

---

All Results Are on a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*)

The assumption or constraint. Every experiment in the paper — search against PRMs, iterative revisions, compute-optimal allocation, and the FLOPs-matched pretraining comparison — uses the MATH benchmark (Hendrycks et al., 2021, specifically the Lightman et al., 2022 split of 12,000 training / 500 test questions) and PaLM 2-S* (Codey) as the base model. The authors state (Section 4):

"We believe this model is representative of the capabilities of many contemporary LLMs"

but provide no cross-model or cross-benchmark validation.

The consequence. Several aspects of the findings could be model-specific or benchmark-specific in ways that affect their generality:

• PRM quality and over-optimization behavior. The paper documents PRM over-optimization as a central limiting factor (Section 5.3, Figure 3 right): beam search degrades easy-problem performance at high budgets. The severity of this over-optimization depends on the PRM's calibration, which in turn depends on the base model's output distribution. A model with different calibration properties (e.g., one that produces more diverse or differently structured solutions) might exhibit qualitatively different difficulty-dependent scaling curves — potentially shifting the difficulty thresholds at which beam search becomes harmful or changing which strategy is optimal per bin.

• Revision model capability. The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. PaLM 2-S* may be particularly amenable to the edit-distance-based pairing procedure described in Section 6.1; other models might require different data construction strategies to learn effective revision.

• MATH-specific structure. MATH consists exclusively of competition-level math problems requiring symbolic reasoning with ground-truth answers checkable via exact string matching (via the grading function from Lightman et al., 2022). It is unclear whether the difficulty-dependent patterns — beam search hurting easy problems, sequential revisions outperforming parallel on easy problems, intermediate sequential-to-parallel ratios being optimal on medium problems — generalize to code generation (where correctness is verified by unit tests), logical reasoning, scientific QA, or tasks requiring factual knowledge rather than multi-step inference. The paper's entire difficulty taxonomy (five quintiles based on pass@1 rate) is defined relative to MATH problems; different benchmarks would produce different difficulty distributions and potentially different optimal strategies.

• The $14\times$ larger model comparison. The FLOPs-matched comparison in Section 7 uses a single larger model from the same family. The finding that test-time compute can substitute for pretraining on easy-to-medium problems (Figure 9) is established for this specific model pair; whether the same holds for other model families (e.g., LLaMA, GPT, Chinchilla-optimal models) or different scale ratios is unknown.

What evidence exists in the paper. The paper contains no cross-model experiments and no cross-benchmark experiments. The 500-question test set, split into five difficulty quintiles of approximately 100 questions each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on approximately 50 questions per fold per bin. No confidence intervals are reported on the compute-optimal scaling curves (Figures 4, 8), so the statistical reliability of the selected strategies at this sample size is unclear. The revision model is trained on the 12,000 MATH training questions — it is unknown whether the revision skill transfers to out-of-distribution math problems or to non-math domains.

Mitigation status. The authors do not claim cross-model or cross-benchmark generality beyond the "representative" assertion quoted above. The limitation is not explicitly discussed as a threat to validity. No future work is suggested regarding broader empirical validation.

---

Hard Problems Remain Essentially Unsolved — Test-Time Compute Cannot Create Capability From Nothing

The assumption or constraint. The compute-optimal framework assumes that the base model has some non-trivial probability of producing a correct solution for a given prompt — the pass@1 rate must be measurably above zero for search or revisions to help. This is implicit in the difficulty estimation procedure (Section 3.2), which computes pass@1 from 2048 samples and bins questions into quintiles. The hardest bin (bin 5) consists of questions where the base model's pass@1 is near zero.

The consequence. Across all methods — search, revisions, and their compute-optimal combinations — the hardest questions show near-zero improvement regardless of compute budget:

• In Figure 3 (right, bottom row), bin 5 accuracy hovers at approximately 1–3% for both best-of-N and beam search at all budget levels (4–256 generations). No method makes meaningful progress.
• In Figure 7 (right, rightmost panel), bin 5 shows approximately 2–3% accuracy irrespective of the sequential-to-parallel ratio at a fixed budget of 128 generations. The curve is essentially flat.
• In the FLOPs-matched comparison (Figure 9, bottommost line in each panel), the bin 5 scaling line is essentially flat near 0–5%, and lies below the $14\times$ larger model's greedy performance across all values of $R$. For PRM search at $R \gg 1$, hard questions show a −52.9% relative disadvantage from using test-time compute instead of the larger model (Figure 1, bottom-right bar chart).

This is a fundamental capability bound: test-time compute can amplify existing capability (by finding correct solutions the model already produces at some low rate, or by refining nearly-correct solutions) but cannot create capability from nothing. If the base model's pass@1 is effectively zero on a problem class, no amount of search or revision will help — there are no correct solutions in the proposal distribution to find or refine.

What evidence exists in the paper. The bin 5 results in Figures 3, 7, and 9 are consistent and unambiguous. The paper is candid about this limitation in the Section 7 takeaway box, noting that for hard problems "pretraining is almost always more effective" and that "test-time compute provides essentially zero benefit regardless of budget, meaning that some capabilities can only be acquired through pretraining, not recovered at inference time." The FLOPs-matched comparison provides the starkest quantification: at $R \gg 1$ with PRM search, hard problems show a −52.9% relative disadvantage.

Mitigation status. This is not a limitation the paper attempts to solve — it is a boundary condition that the paper empirically characterizes and honestly reports. The authors frame test-time compute as "amplifying existing capability" rather than creating new capability, and the bin 5 results are presented as evidence for this boundary. No mitigation is proposed because the limitation is inherent to the approach: if the model never generates correct solutions, no selection or refinement mechanism can produce one. The practical implication is that for genuinely hard or out-of-distribution problems, scaling pretraining (or improving the base model through other means) remains the only viable path.

---

The $14\times$ Larger Model Baseline Is Not Compute-Optimally Trained and Uses No Test-Time Compute of Its Own

The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters by approximately $14\times$ while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The authors acknowledge (Section 7) that this departs from compute-optimal pretraining (Hoffmann et al., 2022), where both data and parameters are scaled equally:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the $14\times$ larger model uses only greedy decoding with no test-time compute augmentation — no majority voting, no best-of-N, no search of any kind.

The consequence. Both choices make the pretraining baseline weaker than it should be for a fair comparison:

• Non-compute-optimal pretraining. A Chinchilla-optimal model trained with $14\times$ more total FLOPs (scaling both parameters and data) would likely outperform a parameter-only-scaled model at the same FLOP budget. The reported advantages of test-time compute over pretraining — e.g., +27.8% on easy questions at $R \ll 1$ (Figure 1, top-right bar chart) — may shrink or reverse against a properly compute-optimal larger model. The comparison answers the question "is test-time compute better than naïvely scaling parameters alone?" but not the more practically relevant question "is test-time compute better than spending the same FLOPs on a properly scaled larger model?"

• No test-time compute for the larger model. Giving the $14\times$ larger model even a modest test-time compute budget — say, best-of-8 or best-of-16 — would create a much stronger baseline. The paper's FLOPs accounting (Section 7) could incorporate this by giving the larger model its own test-time budget: if the total FLOP budget is fixed and the larger model uses more FLOPs per token, its per-query inference budget would be smaller, but it could still apply a lightweight strategy (e.g., best-of-4 with the larger model vs. best-of-256 with the smaller model). The current comparison — smaller model with optimized test-time compute vs. larger model with greedy decoding — conflates the effect of test-time compute with the effect of model scale, making it impossible to attribute the performance difference cleanly to one factor.

What evidence exists in the paper. Figure 9 shows the $14\times$ larger model's greedy performance as stars at three x-axis positions corresponding to three values of $R$. The text in Section 7 acknowledges the non-compute-optimal pretraining choice but does not quantify its impact. The bar charts in Figure 1 report the percentage differences explicitly. No ablation compares against a Chinchilla-optimal larger model or a larger model with its own test-time compute budget.

Mitigation status. The authors explicitly flag compute-optimal pretraining as future work (Section 7 quote above). The greedy decoding choice for the larger model is not discussed as a limitation, though the paper's framework could naturally incorporate it (since the compute-optimal policy could, in principle, be computed for any model, including the larger one). The current comparison should be interpreted as a lower bound on the pretraining advantage: a stronger pretraining baseline would reduce or reverse the reported test-time compute advantages, particularly on medium and hard problems.

---

Revisions and PRM Search Are Never Combined, Leaving Performance on the Table

The assumption or constraint. The paper studies two complementary mechanisms for test-time compute — PRM-guided search (Section 5) and iterative revisions (Section 6) — but evaluates them entirely independently. There is no experiment combining PRM tree-search with the revision model as the proposal distribution. Section 8 explicitly acknowledges:

"we did not experiment with PRM tree-search techniques in combination with revisions"

The consequence. The paper's results represent a lower bound on what a fully integrated system could achieve. The two mechanisms have complementary, difficulty-dependent strengths that suggest natural combinations:

• Revisions improve the proposal distribution (generating better candidates in the first place), while PRM search improves candidate selection (finding the best among generated candidates). Applying beam search over revision model outputs — where each beam step conditions on previous revisions — could yield more accurate step-level assessments because the revision model produces higher-quality candidates than the base model, and the PRM could guide which revisions to pursue rather than blindly generating long revision chains.

• On medium-difficulty problems, beam search outperforms best-of-N (Figure 3, right, bin 3–4) and a balanced sequential-to-parallel ratio outperforms purely sequential or purely parallel revision (Figure 7, right, bins 3–4). Combining beam search with the revision model — using the PRM to score revision steps and prune unpromising revision chains — could break through the performance ceiling that each method individually hits. The paper's compute-optimal policy separately selects the best search strategy and the best revision ratio per difficulty bin, but never jointly optimizes them.

• The paper documents that the base-LM PRM does not transfer well to revision model outputs due to distribution shift (Appendix J, Figure 15a), requiring a separate ORM trained on revision model outputs. A PRM trained directly on revision model trajectories could potentially guide revision search more effectively.

What evidence exists in the paper. The independent evaluations in Sections 5 and 6 show that search and revisions each individually achieve approximately 39–44% accuracy at 256 generations (Figures 4 and 8). The paper does not report any combined results, nor does it estimate the potential gain from combination. The distribution shift finding (Figure 15a) confirms that naive combination (using the base-LM PRM with the revision model) would underperform, but a properly integrated system (with a revision-aware verifier) remains unexplored.

Mitigation status. The authors acknowledge this as a natural next step in Section 8: "we did not experiment with PRM tree-search techniques in combination with revisions, nor distilling the outputs... back into the base LLM." No mitigation is attempted. This is a significant gap because it means the paper does not demonstrate what a fully realized compute-optimal system — simultaneously optimizing over search algorithm, revision depth, and sequential-to-parallel ratio — could achieve.

---

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate, and Revision Training Is Fragile

The assumption or constraint. The revision model is trained solely on sequences where all in-context answers are incorrect, followed by a correct target (Section 6.1). This means the model never sees examples of what to do when the current answer is already correct — it has no signal for when to stop revising or to preserve a correct answer. Section 6.1 reports:

"the model may encounter correct answers in its context (produced during earlier revisions) and incorrectly 'revise' them into wrong answers. The paper reports that approximately 38% of correct answers get converted back to incorrect ones using a naive approach."

Additionally, Appendix K shows that an attempt to further optimize the revision model using ReST$^\text{EM}$ (Singh et al., 2024) degraded performance: at 256 generations, fully sequential performance dropped to approximately 33.5% compared to roughly 38.5% at the optimal ratio for the non-ReST$^\text{EM}$ model (Figure 16). The authors hypothesize that "on-policy data collection in ReST$^\text{EM}$ exacerbates spurious correlations in revision data."

The consequence. The 38% reversion rate means that even when the revision model produces a correct answer mid-chain, there is a substantial probability it will be destroyed in subsequent revision steps. This forces the system to rely on majority voting or verifier-based selection across the entire chain (Section 6.1) to recover correct answers that were generated but then overwritten. These post-hoc selection mechanisms are imperfect patches — they cannot recover a correct answer if all later revisions in the chain are also incorrect, and they introduce additional complexity and potential failure modes. A more principled solution — training the model to recognize when no revision is needed, or using the PRM to decide adaptively whether to continue revising — is not developed.

The ReST$^\text{EM}$ failure is more concerning: it suggests that the revision training procedure is fragile and sensitive to data generation methodology in ways that are not fully understood. The offline data construction procedure (Section 6.1) — pairing independently sampled correct and incorrect solutions post-hoc using edit distance — is carefully designed to teach targeted corrections. Attempting to optimize this with on-policy RL (ReST$^\text{EM}$) breaks the approach entirely, indicating that the positive results depend on specific, potentially brittle choices in data construction. For a practitioner, this means that replicating the revision model's performance requires faithfully reproducing the offline data generation pipeline; modifications or "improvements" (like on-policy fine-tuning) may actively harm performance.

What evidence exists in the paper. The 38% reversion rate is reported in Section 6.1 without a dedicated figure or table — it is a textual claim without detailed breakdown (e.g., does the reversion rate vary by difficulty bin? By position in the revision chain?). The ReST$^\text{EM}$ failure is documented in Appendix K and Figure 16 with one configuration. No ablation studies explore alternative training strategies that might reduce the reversion rate (e.g., including correct-to-correct examples in training, using a stopping criterion, or training a separate "revision-needed" classifier).

Mitigation status. The paper mitigates the reversion problem at inference time via within-chain selection (majority voting or verifier-based selection across all revisions), but this is a post-hoc fix that does not address the root cause (inappropriate training data). The ReST$^\text{EM}$ failure is reported as a negative result without proposed solutions. The authors do not suggest modifications to the training procedure to reduce the reversion rate or improve robustness to training methodology. This limitation is particularly important for practitioners attempting to deploy revision models, as the 38% reversion rate means that nearly two-fifths of the model's successes are self-sabotaging, and the selection mechanisms are only partial remedies.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper does not propose a new algorithm that narrowly beats baselines on a benchmark — it proposes a new category of distributional RL algorithm and, in doing so, reframes what it means to perform distributional reinforcement learning computationally. The shift is from "distributions are the primary objects; sketches are a lossy storage format" to "sketches are the primary objects; distributions are implicit in precomputed Bellman coefficients." This is a conceptual inversion comparable to the kernel trick in machine learning: just as the kernel trick revealed that one can work entirely in inner-product space without explicitly constructing high-dimensional feature maps, the sketch Bellman operator reveals that one can perform distributional Bellman backups entirely in ℝ^m without ever constructing, manipulating, or sampling from probability distributions.

What makes this a reframing rather than an incremental improvement. The paper's core move — defining Bellman coefficients B_r via least-squares regression on φ(r + γG) against φ(G) and using them to define a linear operator 𝒯^π_φ on sketch space — eliminates an entire computational phase (imputation) that prior sketch-based work (Rowland et al., 2019) treated as conceptually necessary. Figure 1 visualizes this cleanly: the SFDP pathway has three stages (imputation → distributional Bellman backup → sketch extraction), while Sketch-DP collapses them into a single linear update involving only precomputed matrix-vector products. This is not optimizing the imputation step — it is recognizing that the imputation step was never required in the first place, provided one is willing to accept the approximation error from the regression-based Bellman coefficients. The error propagation theory in Section 4 then shows this approximation is controllable: it shrinks as O(1/m) for appropriately designed sketches (Proposition 4.4), meaning arbitrary accuracy is achievable, not just heuristic improvement.

Resolving prior contradictions. The paper reconciles a tension that had been treated as a binary classification in the literature. Rowland et al. (2019, Theorem 4.3) characterized all Bellman-closed mean embedding sketches — they are limited to linear combinations of the first m moments. This implied that any sketch beyond low-order moments was necessarily incompatible with direct sketch-space Bellman backups, forcing the expensive and biologically implausible imputation round-trip. Meanwhile, moment-based methods (Sobel, 1982; Tamar et al., 2013; 2016) were known to be numerically unstable at scale due to the wildly diverging magnitudes of higher moments (Appendix D.1, Figure 7). The field was stuck between two unattractive options: exact DP with numerically unstable moments, or imputation-based approximate DP with high computational cost and no convergence theory.

The paper cuts through this dilemma by showing that exact Bellman closure is unnecessarily strict. By treating the closure condition as a regression target — find B_r minimizing 𝔼_G∼μ[‖φ(r + γG) − Bφ(G)‖²] rather than requiring pointwise equality — the framework achieves approximate closure for virtually any feature map, with the regression residual directly bounding the per-step error (Proposition 4.1). This converts the design of distributional RL algorithms from a classification problem ("is this sketch Bellman-closed?") into an approximation problem ("how well can B_r fit the bootstrap transformation under μ?"). The practical consequence is that feature maps with excellent numerical properties — bounded, localized sigmoids and Gaussians — become usable for direct sketch-space DP and TD, escaping the moment straightjacket while retaining the speed and theoretical tractability of linear updates.

Which directions become more attractive. Several research threads gain new motivation from this work:

• Feature map design for distributional RL becomes a first-class research question rather than an afterthought. The translation family (Equation 8) provides a starting point, but the framework invites systematic exploration of learned features, multi-scale features, adaptive anchor placement, and domain-specific feature maps tailored to capture risk-sensitive quantities (tail expectiles, quantiles, variance). The paper's empirical sweeps over κ, m, and s (Figures 4, 9) demonstrate that feature choice matters substantially for both mean embedding accuracy and distributional recovery, but do not exhaust the design space.

• Neuroscience of distributional learning gains a biologically plausible computational mechanism. Tano et al. (2020) criticized imputation-based sketch methods as requiring neurons to explicitly reconstruct probability distributions — a level of explicit probabilistic computation unsupported by neural evidence. The sketch Bellman operator U(x) ← 𝔼[B_R U(X')] is a purely local, linear propagation rule that could plausibly be implemented by synaptic weight matrices encoding B_r. The non-normal dynamics observed in Appendix B.5 (transient expansion under single B_r multiplication, stable convergence under iteration) are consistent with dynamical properties observed in cortical circuits (Hennequin et al., 2012; Bondanelli & Ostojic, 2020), providing a candidate mechanism for how neural populations might perform distributional RL without explicit probabilistic inference.

• Linear operator theory in RL analysis becomes newly applicable to distributional methods. The sketch Bellman operator 𝒯^π_φ is linear — a property not shared by categorical projection, quantile regression, or imputation-based sketch methods. This opens the door to importing spectral analysis, perturbation theory, and linear dynamical systems theory into distributional RL analysis, as the paper partially demonstrates through its singular value and eigenvalue analysis in Appendix B.5.

Which directions become less central. The paper's framework implicitly argues against continued investment in imputation-based sketch methods (SFDP/SFTD). While expectile-based sketches drove important advances (Dabney et al., 2020; Lowet et al., 2020), the empirical comparison in Figure 11 shows Sketch-DP achieving better accuracy than SFDP while running >100× faster. The conceptual advance — that imputation is unnecessary — suggests that future work on sketch-based distributional RL should build on the Bellman coefficient framework rather than designing more sophisticated imputation strategies. Similarly, the failure of moment-based TD learning at m = 50 (Figure 7) — performing worse than the Dirac baseline — suggests that exact moment DP, while theoretically elegant (Sobel, 1982), is a dead end for scalable distributional RL, and that approximate closure via regression is the path forward.

Follow-Up Research This Work Enables

Convergence theory for Sketch-TD under linear function approximation. The paper provides rigorous convergence analysis for the DP case (Section 4) but explicitly leaves TD convergence as future work. The natural next step is to analyze the stochastic approximation U_{k+1}(x_k) ← (1 − α_k)U_k(x_k) + α_k B_{r_k} U_k(x'_k) under standard assumptions (diminishing learning rates, ergodic Markov chain sampling, linear function approximation U_θ(x) = Φ(x)θ). The DP analysis template — decomposing error into Bellman approximation, reconstruction, and embedding components — should extend to the TD case via Kushner-Yin ODE method (Kushner & Yin, 1997), but the stochastic setting introduces Martingale noise and off-policy sampling corrections that do not appear in the DP analysis. A strong result would establish that the TD iterates converge to a neighborhood of the true sketch values with radius proportional to the Bellman approximation error ε_B plus an additional term from temporal-difference bias, and characterize the asymptotic variance. The theory could be tested empirically by measuring ‖U_TD − U^π‖ as a function of training steps and comparing against the predicted asymptotic bound, using the same tabular MRPs as Figure 7 but with linear function approximation over random features.

Learned Bellman coefficients via meta-gradient or online adaptation. The paper's framework requires precomputing B_r from a fixed weighting distribution μ over an estimated return range. This is practical when the return range is known and rewards are discrete (as in Atari with clipped rewards in {−1, 0, 1}), but breaks down when the return distribution shifts during training (e.g., policy improvement in Q-learning), when the reward set is large or continuous, or when the return range is poorly estimated. An important extension would be to learn B_r from data, either by meta-gradient (differentiating through the Bellman coefficients with respect to the downstream RL objective) or by online adaptatation (updating B_r from observed returns during training, treating the regression problem as a streaming least-squares problem with exponential forgetting). The key challenge is maintaining the linear structure — if B_r is learned online, it must remain a linear operator to preserve the framework's properties, but the data distribution shifts as the policy improves. A concrete experiment would compare fixed precomputed B_r against online-adapted B_r on an Atari game with unclipped rewards (where the return range expands during training), measuring both sketch accuracy and downstream agent performance.

Combining sketch-based distributional RL with risk-sensitive policy optimization. The framework's ability to target specific distributional features through the choice of φ is currently used only for policy evaluation, with policy improvement driven by the expected value readout β (Section 5.1). However, the learned sketch values contain information about variance, tail risk, and other distributional properties that the value readout discards. A natural extension is risk-sensitive policy improvement where the greedy action is selected based on a risk-adjusted utility computed from the sketch, e.g., a^⋆ = arg max_a ⟨β, U(x, a)⟩ − λ · variance_estimate(x, a), where variance_estimate is another linear readout from the same sketch values (precomputed via a regression analogous to β but predicting squared returns). The framework makes this straightforward because once U(x, a) is learned, any distributional functional expressible as 𝔼[f(G)] for some f in the span of φ can be read out linearly. A concrete experiment on a risk-sensitive gridworld (e.g., a cliff-walking variant with catastrophic low-probability negative outcomes) would test whether sketch-based risk-sensitive policies outperform expected-value policies and quantile-based risk-sensitive methods (Dabney et al., 2018b).

Multi-scale and hierarchical feature maps for long-horizon distributional RL. The O(1/m) convergence rate in Proposition 4.4 requires m to grow with 1/(1 − γ) to maintain fixed accuracy — longer horizons require more features because return distributions become more complex (wider support, more multi-modality from stochastic transitions). For γ ≈ 0.99 (common in continuous control), the required m may be impractically large with uniform anchor grids. A promising direction is multi-scale feature maps with non-uniform anchor spacing: dense anchors in high-probability regions of the return distribution (near the mean) and sparse anchors in the tails, or features with varying slopes to capture both coarse and fine distributional structure. The translation family (Equation 8) could be extended to ϕ_i(z) = κ(s_i(z − z_i)) where both s_i and z_i are optimized rather than uniform. The extension is natural because the Bellman coefficient framework only requires that the features be linearly independent under μ — it does not assume uniform spacing or identical slopes. A concrete experiment on a chain MDP with γ = 0.99 would measure the number of uniform-grid features needed to achieve a target Cramér distance, then test whether optimized non-uniform features achieve the same accuracy with fewer features.

Empirical characterization of the over-optimization phase transition in distributional RL. The paper's analysis of B_r's spectral properties (Appendix B.5) reveals that the Bellman coefficients are generally non-normal: their singular values can exceed 1 (single-step expansion) while eigenvalues have magnitude ≤ 1 (long-term stability). This suggests a phase transition in Sketch-DP dynamics: for small m (poor approximation), the error propagation may diverge; for sufficiently large m, the contraction of 𝒯^π in Wasserstein distance dominates and the iterates converge to a neighborhood of U^π. The paper's experiments (Figures 4, 9) show monotonic improvement with m, but do not probe the low-m regime where divergence might occur. A systematic study sweeping m from 2 to 200 on environments with varying γ and stochasticity, measuring whether the DP iterates converge or diverge, would map out the stability boundary and test whether the theoretical bound in Proposition 4.3 is tight — i.e., whether the predicted error floor matches the empirically observed asymptotic error. This would also characterize what the paper only hints at: there is a minimum m below which the sketch Bellman operator is not just inaccurate but unstable, and this minimum grows with the horizon 1/(1 − γ) and the stochasticity of the environment.

Anchoring the framework in the kernel methods literature via explicit kernel choice. The mean embedding sketch with feature map ϕ corresponds to a reproducing kernel K(z, z') = ⟨ϕ(z), ϕ(z')⟩. The paper treats ϕ as the primary design choice, but the kernel perspective may provide more principled guidance: rather than hand-designing ϕ, choose a kernel K with known properties (universality, characteristic-ness, decay rate) and use random Fourier features, Nyström approximation, or Cholesky features to obtain finite-dimensional ϕ. The connection is that mean embeddings in finite-dimensional Euclidean space approximate mean embeddings in the RKHS ℋ_K; the quality of this approximation is controlled by standard kernel approximation theory (Sriperumbudur et al., 2010). A concrete follow-up would instantiate the sketch framework with random Fourier features for a Gaussian kernel, compare against hand-designed translation-family features at equal m, and test whether the kernel-derived features achieve lower Cramér distance for the same dimensionality — providing a principled feature selection strategy that the current heuristics (Appendix B.6) lack.

Practical Applications and Downstream Use Cases

Faster distributional RL training in research and industry. The per-iteration speedup of >100× over SFDP (Figure 11) and the competitive frame rate of Sketch-DQN vs. QR-DQN (Table 1: 1326 vs. 1258 fps) mean that the sketch Bellman operator is immediately actionable for any practitioner currently using distributional RL. For research labs running large-scale distributional RL experiments (e.g., multi-seed Atari sweeps, continuous control with distributional critics), replacing QR-DQN or C51 with Sketch-DQN provides comparable or better median performance (Figure 5: Sketch-DQN median ~1.7 vs. QR-DQN ~1.4) at slightly higher throughput, with the added benefit that the same learned sketch values can be mined for multiple distributional statistics (variance, tail expectiles) without retraining. The precomputation of B_r is a one-time O(m³) cost that amortizes over millions of training steps — for m = 401 as in the Atari experiments, the C matrix is 401×401, which inverts in <1 second on modern hardware. The practical barrier to adoption is the feature map design: the heuristics in Appendix B.6 (anchors slightly wider than return range, slope chosen for 50% overlap of non-trivial support) provide a starting point, but practitioners should budget tuning time for m, s, and κ on a held-out environment subset, as the paper shows meaningful sensitivity (Figure 13: sigmoid outperforms Gaussian by ~1–2 mean normalized return points).

Targeted risk assessment from learned sketches without full distribution reconstruction. The framework enables a workflow that is awkward in categorical or quantile methods: train a Sketch-DQN once, then read out multiple distributional statistics from the same learned sketch values by precomputing different readout coefficients β_f for different functionals f (e.g., f(G) = G for the mean, f(G) = (G − β_mean · ϕ(G))² for variance, f(G) = 𝟙{G ≤ τ} for the CDF at threshold τ). Because the readout is a linear function of the sketch, evaluating a new functional requires only a dot product, not retraining. Concretely, a financial RL application might use the same sketch to compute both the expected portfolio return (for allocation decisions) and the 5th percentile return or expected shortfall (for regulatory risk reporting). The accuracy of the risk readout depends on whether the distributional feature of interest is well-captured by the chosen ϕ — sigmoid features with uniform anchors will capture tail quantiles less accurately than the mean, and a risk-focused deployment would want anchor points concentrated in the left tail. The paper's Figure 10 quantifies this sensitivity: anchors narrower than the return range cause sharp Cramér distance degradation, providing actionable guidance for anchor placement.

Biologically-inspired neural network architectures for distributional learning. The paper's finding that the sketch Bellman operator is linear, local, and exhibits non-normal dynamics (Appendix B.5) provides a concrete computational hypothesis for how biological neural circuits might implement distributional RL — a question of active interest in computational neuroscience (Dabney et al., 2020; Lowet et al., 2020; Tano et al., 2020). A neuromorphic hardware implementation or spiking neural network could implement the sketch update U(x) ← 𝔼[B_R U(X')] using a recurrent weight matrix encoding the reward-averaged Bellman coefficients Ḃ, with the non-normal dynamics (transient amplification followed by stable convergence) matching observed neural response patterns (Hennequin et al., 2012). This is speculative but testable: one could train a recurrent neural network on a distributional RL task, extract the effective linear dynamics around the fixed point, and check whether they match the structure predicted by Bellman coefficients — specifically, singular values > 1 with eigenvalues < 1. The paper's open-source code and the explicit form of B_r for Gaussian features under Lebesgue measure μ (Appendix B.3) provide a computational target for such an analysis.