ArXiv: 2402.08530
🎯 Pitch
This paper introduces the distributional successor measure, which separates transition dynamics from reward in distributional RL—enabling zero-shot evaluation of novel reward functions. The resulting δ-model uniquely predicts full return distributions without retraining, correctly ranking policies by both mean return and risk-sensitive criteria across tasks—a capability no previous method achieves.
1. Executive Summary
This paper introduces a new approach for distributional reinforcement learning that cleanly separates transition structure from reward information, enabling zero-shot distributional policy evaluation on novel tasks without further learning. The key contribution is the distributional successor measure (a distribution over state-occupancy distributions that captures aleatoric uncertainty about an agent's future trajectories) and a tractable approximation called the δ-model (an equally-weighted particle representation where each atom is itself a generative model of state, trained via a two-level maximum mean discrepancy loss). On continuous control domains—Windy Gridworld and Pendulum—the δ-model substantially outperforms baselines at zero-shot return distribution prediction while also enabling risk-sensitive policy selection across held-out reward functions, establishing that policies can be correctly ranked by both mean return and conditional value-at-risk without any additional data collection or training—a capability that no prior method, including standard successor measures or learned transition models, can provide.
2. Context and Motivation
The Core Problem: Reward Functions Are Often Unknown Until After Training
The fundamental problem this paper addresses is a structural limitation in how distributional reinforcement learning handles reward information. In standard distributional RL—pioneered by Morimura et al. (2010) and Bellemare et al. (2017a)—an agent learns the entire probability distribution of returns (cumulative discounted rewards) for a given policy. This distribution captures not just the expected performance but also the variability of outcomes, enabling risk-aware decision-making where an agent might prefer a policy with lower variance even if its mean return is slightly worse (Dabney et al., 2018a; Zhang & Weng, 2021; Fawzi et al., 2022).
However, there is a critical practical constraint that the paper's introduction highlights directly:
"A notable drawback of existing approaches to distributional RL is that rewards must be available at training time in order to predict the return distribution. For example, if we wish to evaluate a trained policy on a new task with regard to various performance criteria, these predictions of the return distributions must be trained from scratch."
This means that if you have trained a policy for one objective (say, navigating a robot to a target location while minimizing energy consumption), and you later want to evaluate how that same policy performs under a different reward function (say, maximizing speed while avoiding certain danger zones), standard distributional RL provides no mechanism for doing so without collecting new data and retraining. The return distribution is fundamentally tied to the specific reward function used during training, and there is no clean factorization that separates "what the agent does" from "what the agent gets rewarded for."
This limitation matters for several concrete reasons:
-
Safety-critical deployment: Before deploying a learned policy in a real system, engineers often need to evaluate it against multiple risk criteria—value-at-risk, conditional value-at-risk, probability of catastrophic failure—that correspond to different implicit reward structures. Standard distributional RL requires retraining for each such evaluation, which may be expensive or dangerous if it requires new data collection from the actual system.
-
Multi-objective decision-making: In many applications, the true reward function is not known precisely at training time but emerges from evolving stakeholder preferences or environmental constraints. An agent that can produce return distributions for arbitrary reward functions post-hoc enables flexible, on-the-fly tradeoff analysis without additional learning.
-
Reusing learned behaviors: If a policy encodes sophisticated motor skills or navigation strategies, it is wasteful to require full retraining simply because the evaluation criterion changed. A method that separates the policy's behavioral consequences from reward evaluation would allow the same trained policy to be assessed against many objectives.
The Successor Representation Solves This for Expected Returns—But Not for Distributions
The paper builds directly on a known solution to this factorization problem for expected returns. The successor representation (SR; Dayan, 1993) and its continuous-state generalization, the successor measure (SM; Blier et al., 2021; Touati & Ollivier, 2021), provide exactly this separation: they learn the discounted expected visitation frequencies of states under a policy, independent of any reward function. Equation 2 in the paper captures the core insight:
In words: the value function (expected return) for any reward function can be computed as an expectation of under the successor measure , which itself depends only on the policy and environment dynamics, not on rewards. This enables zero-shot policy evaluation: learn once, and then compute for any new reward function instantly, without further learning or data collection. Recent work has shown this can be done at scale in continuous domains using generative models (Janner et al., 2020; Thakoor et al., 2022) or density models (Blier et al., 2021; Blier, 2022).
But the successor measure only captures expected visitation frequencies. It answers questions like "what fraction of time does the agent spend in dangerous states, on average?" but cannot answer questions like "what is the distribution over the fraction of time spent in dangerous states across different possible trajectories?" This distinction—between the expectation and the full distribution—is precisely where distributional RL adds value over classical RL. The successor measure factorizes return expectations, but says nothing about return variability.
Prior Attempts at Distributional Transfer Learning Fall Short
The paper acknowledges a few lines of prior work that attempt to bring distributional thinking to transfer or successor-based methods, and explains why they are insufficient.
Distributional successor features (Gimelfarb et al., 2021; Carvalho et al., 2023) learn categorical distributions over scalar features, applying the distributional RL techniques of Bellemare et al. (2017a) to the successor features framework (Barreto et al., 2017). The critical limitation is that these methods model the marginal distribution of each feature independently—they do not account for the joint distribution over features. The paper states this explicitly:
"Gimelfarb et al. (2021) and Carvalho et al. (2023)... applied the distributional RL techniques of Bellemare et al. (2017a) and Achab et al. (2023) to learn categorical distributions over features, but did not account for their joint distribution."
This is a significant gap because state visitations are not independent across states: if the agent spends a lot of time in state A on a particular trajectory, it necessarily spends less time in state B. The joint distribution captures these dependencies, and without it, the predicted return distribution (which is a linear combination of per-state rewards weighted by visitation frequencies) will be incorrect whenever rewards are correlated across states. Gimelfarb et al. (2021) uses these marginal distributions to optimize an entropic risk objective, but the method cannot reconstruct the full return distribution for an arbitrary reward function.
Multivariate distributional RL (Freirich et al., 2019; Zhang et al., 2021b) does learn joint distributions over finite-dimensional feature vectors. In fact, Zhang et al. (2021b) use an MMD loss for learning multivariate return distributions, which the current paper builds on algorithmically. However, these methods still operate in a reward-dependent regime: they learn the joint distribution of per-step rewards along a trajectory, not the distribution of state occupancies. The reward function must be known at training time. The current paper notes that the multi-dimensional reward distributional Bellman equation studied in this prior work:
"...can be framed as an instance of Equation 5 when X is finite."
but this is exactly the limitation—X being finite means a tabular state space, and the approach collapses to reward-dependent modeling rather than the reward-agnostic occupancy modeling that enables zero-shot transfer.
The tabular "distributional SR" of Vértes & Sahani (2019) considers learning the successor representation in partially observable MDPs (POMDPs) and calls this a "distributional SR," but the term means something different there: it refers to the fact that the agent's belief state is a distribution over the underlying MDP state, not to a distribution over occupancy measures. This is a different sense of "distributional" and does not enable zero-shot return distribution evaluation.
The Gap: No Method Separates Transition Uncertainty from Reward for Full Return Distributions
The paper's central diagnosis is that while the successor measure factors expected returns as , there is no analogous factorization for return distributions. The paper asks: can we define an object (a random successor measure) such that:
holds in distribution for any deterministic reward function ? If so, then learning the distribution of would enable zero-shot prediction of the full return distribution for any reward function—not just its expectation—by plugging into the right-hand side and propagating the randomness of through the expectation.
This would combine the two desirable properties that currently exist only in isolation: the zero-shot capability of successor measures (transfer across reward functions) and the uncertainty quantification of distributional RL (transfer across risk criteria). No existing method provides both.
The Alternative—Learning the Transition Model—Is Theoretically Sufficient but Practically Inadequate
The paper acknowledges a subtle point that the careful reader might raise: since the distribution of trajectories is fully determined by the policy and the transition kernel , one could in principle learn , sample rollouts from it for any reward function, and compute empirical return distributions. This would achieve zero-shot distributional evaluation without the new machinery the paper introduces.
The paper addresses this directly in Section 3, and the experiments confirm the practical failure of this approach:
"Despite this, recovering the distributional SM or the SM from in large MDPs is intractable, and SM-like models are known to be more robust to estimation error for long-horizon prediction in continuous MDPs (Janner et al., 2020; Thakoor et al., 2022; Touati et al., 2023). Crucially, unlike an approach to zero-shot distributional evaluation that estimates return distributions by sampling rollouts from a learned and computing MC returns, the distributional SM is not prone to accumulation of model error, which results in substantially more accurate estimation, as we show in Section 6."
This is a crucial motivation. Learning an accurate one-step transition model is hard enough in continuous state spaces, but the problem compounds dramatically when using that model for multi-step rollouts. Each step's prediction error accumulates, and for the long horizons typical under discount factors like or , the effective horizon is 20 or 100 steps respectively. Even small per-step errors compound to make the terminal states of the rollout essentially meaningless—a phenomenon well-documented in model-based RL (Jafferjee et al., 2020; Abbas et al., 2020; Lambert et al., 2022).
The successor measure and its distributional analogue sidestep this compounding error problem by directly learning a temporally-abstracted representation. Rather than learning to predict one step ahead and chaining predictions 100 times, they learn to predict the discounted sum of future occupancies in one shot—a form of "temporal abstraction" that is inherently more robust.
The Bayesian Analogy: Epistemic vs. Aleatoric Uncertainty
The paper provides a conceptual distinction in Section 6 that clarifies what the distributional SM adds over an ensemble of standard successor measures (like an ensemble of -models; Janner et al., 2020):
"Conceptually, the γ-model ensemble is expected to capture the epistemic uncertainty over the SM, while the distributional SM estimates the aleatoric uncertainty due to randomness of the MDP dynamics and the policy."
An ensemble of -models, each trained independently on different random seeds or data subsets, captures epistemic uncertainty: our ignorance about what the true successor measure is, due to finite data. As we collect more data, all ensemble members should converge to the same , and the ensemble distribution collapses to a point mass.
The distributional SM captures something fundamentally different: aleatoric uncertainty—the inherent randomness in which future states the agent will actually visit, even given perfect knowledge of the environment and policy. Even if we know exactly, a stochastic policy in a stochastic environment produces different trajectories on different runs. The distributional SM models this irreducible randomness as a distribution over occupancy measures. This distinction matters for risk-sensitive evaluation: epistemic uncertainty can be reduced with more data, but aleatoric uncertainty is a property of the policy-environment interaction that must be accounted for regardless of data quantity. A risk-averse decision-maker needs to know the aleatoric variability to properly assess policies; epistemic uncertainty estimates alone would incorrectly suggest that the policy becomes deterministic with enough data.
How This Paper Positions Itself
The paper positions its contribution as filling a conceptual hole in the intersection of successor representations and distributional RL. It does not propose a new reinforcement learning algorithm for policy optimization, nor does it claim to improve sample efficiency or asymptotic performance on standard benchmarks. Instead, it introduces a new mathematical object—the distributional successor measure—and demonstrates that:
-
This object exists and has clean theoretical properties (a distributional Bellman equation, a contractive operator, convergent dynamic programming; Propositions 3.5, 3.6, Corollary 3.7).
-
It can be tractably approximated with a novel architecture (the -model) using techniques from generative modeling and kernel methods, with specific practical innovations (-step bootstrapping, adversarial kernels) that are necessary for stable training (Sections 4 and 5).
-
It enables capabilities—zero-shot distributional policy evaluation and risk-sensitive policy comparison on novel reward functions—that no prior method can provide (Section 6).
The paper explicitly acknowledges that it does not combine the distributional SM with policy optimization or with the successor features framework for generalized policy improvement (Barreto et al., 2017). These are left as natural extensions. The contribution is foundational: defining the right object and showing it can be learned, leaving its integration into full agent architectures for future work.
3. Technical Approach
3.1 Reader Orientation
This paper develops a new mathematical object and learning algorithm for distributional reinforcement learning. The system being built is a learned model—called a δ-model—that, given a starting state and a policy, can generate possible future state-occupancy distributions that the agent might experience, and then use those to compute the full return distribution for any reward function without further training. The problem it solves is that standard distributional RL ties reward information to training time, preventing zero-shot evaluation on new tasks; the δ-model solves this by learning a reward-agnostic distribution over occupancy measures from which return distributions for arbitrary rewards can be derived post-hoc through a simple plug-in operation.
3.2 Big-Picture Architecture (Diagram in Words)
The system has five major components connected in a learning-and-inference pipeline:
-
Random Occupancy Measure (
$M^\pi$) — the mathematical target: a random variable whose realizations are probability distributions over the state space, describing the discounted visitation pattern along a sampled trajectory. This exists in the environment but must be learned from data. -
Distributional Successor Measure (
$\mathcal{M}^\pi$) — the formal object to be approximated: the probability distribution of$M^\pi$, i.e., a distribution over occupancy distributions. For each state$x$,$\mathcal{M}^\pi(x)$is a distribution over$\mathcal{P}(\mathcal{X})$(the space of probability measures on states). -
δ-Model — the tractable approximation to
$\mathcal{M}^\pi$: an equally-weighted particle representation with$m$particles (Dirac masses), where each particle$\theta_i(x)$is itself a generative model of state (a distribution over$\mathcal{X}$). The ensemble of$m$such generators, treated as equally-likely atoms, forms the approximation$\hat{\mathcal{M}}(x) = \frac{1}{m}\sum_{i=1}^m \delta_{\theta_i(x)}$. -
Two-Level MMD Loss — the learning signal: a distributional Bellman residual computed using maximum mean discrepancy at two nested levels. A state-level MMD compares individual generative models
$\theta_i$to each other; a model-level MMD compares the distribution$\hat{\mathcal{M}}(x)$to its bootstrapped target$\tilde{\mathcal{M}}(x)$using a kernel defined on the space of distributions. -
Adversarial Kernel with Invertible Network — the adaptive similarity measure: a learned feature map
$f$(implemented as an iResNet) that transforms raw states before computing the state-level MMD, enabling the kernel to adapt as the generative models evolve during training.
Information flows as follows during training: a state $x$ is sampled from the policy's stationary distribution → a trajectory $(x, x_1, ..., x_n)$ is collected under $\pi$ → the current δ-model generates $m$ occupancy distribution samples $\theta_i(x)$ by conditioning each generator network on $x$ and noise → target occupancy distributions $\bar{\theta}_i(x)$ are constructed by mixing trajectory states with bootstrapped samples from $\theta_i(x_n)$ according to the geometric weighting → the adversarial kernel transforms state samples from all generators → the two-level MMD is computed between the source and target model distributions → the generators are updated to minimize this MMD (making predicted occupancies match the Bellman target), while the kernel network is updated to maximize it (making the distance more discriminative). At inference time for a new reward function: sample model atoms from the δ-model, compute the expected reward under each sampled occupancy distribution using Equation 4, and collect the resulting scalar values into an empirical return distribution.
3.3 Roadmap for the Deep Dive
-
First, the mathematical object: define the random occupancy measure
$M^\pi$and the distributional SM$\mathcal{M}^\pi$, establishing the factorization identity (Proposition 3.2) that is the paper's central theoretical contribution. This gives us the target we want to learn and explains why it enables zero-shot distributional evaluation. -
Second, the distributional Bellman equation for
$M^\pi$(Proposition 3.5) and the contractive operator$\mathcal{T}^\pi$on distributional SMs (Proposition 3.6, Corollary 3.7). These provide the theoretical foundation for temporal-difference learning—showing that$\mathcal{M}^\pi$is the unique fixed point of a$\gamma$-contractive operator, so iterative updates will converge to it. -
Third, the δ-model representation: the equally-weighted particle approximation, why particles live in
$\mathcal{P}(\mathcal{X})$(not in a Euclidean space), and how generative models (like GAN generators) serve as atoms. This is the bridge from theory to implementation—what we actually store and update in code. -
Fourth, the two-level MMD loss construction, working bottom-up: state kernel
$\kappa$for comparing individual state distributions → model kernel$k_\kappa$for comparing distributions-over-distributions → the full Bellman residual loss. This is the most technically involved component and requires careful exposition of why nested MMD is the right choice. -
Fifth, the practical training techniques:
$n$-step bootstrapping (reducing reliance on learned targets for long-horizon predictions) and the adversarial kernel with iResNet (handling non-stationarity of the distributions being modeled). These are the "tricks that make it work" and are independently valuable.
3.4 Detailed, Sentence-Based Technical Breakdown
This is a foundational methods paper whose core idea is that return distributions can be factorized into a reward-agnostic distribution over state occupancies (the distributional SM) and the reward function itself, enabling zero-shot distributional policy evaluation through a plug-in estimator.
The Random Occupancy Measure and Distributional SM
The paper begins by reframing the standard successor measure $\Psi^\pi$ in a way that reveals its "distributional" structure. For a fixed policy $\pi$, the normalized successor measure for an initial state $x$ is:
where $S \subseteq \mathcal{X}$ is any measurable set of states, $\gamma \in [0,1)$ is the discount factor, and $\Pr(X_t \in S \mid X_0 = x)$ is the probability that the agent is in $S$ at time $t$ after starting at $x$ and following $\pi$.
What it computes: this is the expected discounted visitation frequency to states in $S$. The factor $(1-\gamma)$ normalizes $\Psi^\pi(\cdot \mid x)$ to be a probability distribution (it integrates to 1 over $\mathcal{X}$), making it a proper probability measure over states for each starting state $x$. Without this normalization, the successor measure would have total mass $(1-\gamma)^{-1}$, corresponding to the effective horizon.
Why this form: the normalization allows us to treat $\Psi^\pi(\cdot \mid x)$ as a distribution we can sample from—a key property that Janner et al. (2020) exploited to build generative successor models. The $(1-\gamma)\gamma^t$ weights implement a geometric weighting that prioritizes near-term states (visited at small $t$) over far-future states, matching how discounting works in the return.
The paper then rewrites this as an expectation over random trajectories. Let $(X_t)_{t \geq 0}$ be the random sequence of states generated by executing $\pi$ from $X_0 = x$. Then $\delta_{X_t}$ is the Dirac measure at the random state $X_t$ (a random probability distribution that places all its mass on wherever the agent happens to be at time $t$). The successor measure can be expressed as:
What it computes: the expected value (over random trajectories) of a random discounted sum of Dirac measures. Inside the expectation is a sum over time of the indicator $\delta_{X_k}(S) = \mathbf{1}\{X_k \in S\}$, weighted by $(1-\gamma)\gamma^k$. The expectation over trajectories then yields the probability that a state visited at a random discounted time falls in $S$.
Why this form: it reveals that the successor measure is the expectation of a random object—specifically, a random probability distribution. This suggests that if we remove the expectation, we obtain a richer object that captures the full distribution of what the discounted occupancy pattern could be, not just its average.
This motivates Definition 3.1, the random occupancy measure $M^\pi$:
where $X_k$ are the random states generated by $\pi$, $\delta_{X_k}$ is the Dirac measure at the random state $X_k$ (a random probability distribution), and $S \subseteq \mathcal{X}$ is any measurable set.
What it computes: for each starting state $x$, $M^\pi(\cdot \mid x)$ is a random probability distribution over $\mathcal{X}$. Each realized trajectory $(x_0, x_1, x_2, ...)$ produces one specific occupancy distribution—a weighted sum of Dirac masses at the visited states. If the agent visits state A three times early in the trajectory and state B once late, the occupancy distribution will put more mass on A than B (due to discounting). Across different trajectories, the occupancy distribution varies, capturing the aleatoric uncertainty in where the agent actually goes.
Why this form: this is the object that cleanly separates transition structure from reward. The key insight is that $M^\pi$ depends only on the policy $\pi$ and the environment transition dynamics $p^\pi$—it is entirely reward-agnostic. As we will see, any reward function can be combined with $M^\pi$ to obtain the corresponding return distribution.
Critical detail about the nature of $M^\pi$: for each starting state $x$, $M^\pi(\cdot \mid x)$ is a random variable taking values in $\mathcal{P}(\mathcal{X})$ (the space of probability measures over states). Its distribution is therefore a distribution over probability distributions—a point on $\mathcal{P}(\mathcal{P}(\mathcal{X}))$. The paper emphasizes this with explicit notation in the prose:
"So, for any measurable
$Y \subset X$,$M^\pi(Y \mid x)$is also a random variable, which gives the discounted proportion of time spent in$Y$across different possible sampled trajectories. Thus, the distribution of$M^\pi(\cdot \mid x)$is a distribution over probability distributions; see Figure 1."
Figure 1 (left) illustrates this with a T-maze example: a policy that takes different branches with different probabilities produces three possible occupancy distributions $\theta_1, \theta_2, \theta_3$ (each corresponding to a distinct behavioral outcome), and the distributional SM is the distribution $\frac{1}{6}\delta_{\theta_1} + \frac{1}{2}\delta_{\theta_2} + \frac{1}{3}\delta_{\theta_3}$ over these atoms.
Definition 3.4 formalizes the distributional successor measure (distributional SM) $\mathcal{M}^\pi$:
where $\text{Law}(M^\pi(\cdot \mid x))$ denotes the probability distribution of the random occupancy measure $M^\pi(\cdot \mid x)$, and $\mathcal{M}^\pi: \mathcal{X} \to \mathcal{P}(\mathcal{P}(\mathcal{X}))$ is a function mapping each state to a distribution over state distributions.
What it computes: for each state $x$, $\mathcal{M}^\pi(x)$ is a probability measure on $\mathcal{P}(\mathcal{X})$—it tells us which occupancy distributions are likely and which are unlikely, given the policy's stochastic behavior and the environment's stochastic dynamics.
Why this form: $\mathcal{M}^\pi(x)$ is the reward-agnostic object that we aim to learn. Once we have it, we can recover return distributions for any deterministic reward function via Proposition 3.2.
Proposition 3.2 is the paper's central theoretical result—the factorization identity:
where $G^\pi_r(x) = \sum_{t=0}^\infty \gamma^t r(X_t)$ is the random discounted return (starting at $X_0 = x$), $r: \mathcal{X} \to \mathbb{R}$ is any deterministic reward function, and $\stackrel{\mathcal{L}}{=}$ means equality in distribution (the two random variables have the same probability law).
What it computes: the random return $G^\pi_r(x)$ (a scalar random variable) has the same distribution as $(1-\gamma)^{-1}$ times the expectation of the reward function $r$ under the random occupancy measure $M^\pi(\cdot \mid x)$. The right-hand side is random because $M^\pi(\cdot \mid x)$ itself is a random distribution—each realized trajectory produces a different occupancy distribution, and the expectation of $r$ under that occupancy distribution is a different scalar. Passing these scalars through the $(1-\gamma)^{-1}$ scaling yields exactly the return distribution.
Why this form: it cleanly factorizes the return distribution into two independent components: (1) the distributional SM $\mathcal{M}^\pi$, which captures all transition and policy randomness and is independent of rewards; and (2) the reward function $r$, which captures the task-specific evaluation criterion. This factorization is what enables zero-shot distributional policy evaluation: learn $\mathcal{M}^\pi$ once from interaction data (no rewards needed), and then for any new reward function $r$, simply compute the right-hand side to obtain the corresponding return distribution without any additional data collection or learning.
Why the reward function must be deterministic: the paper notes in Remark 3.3 and expands in Appendix C that this factorization fails for stochastic rewards due to the statistical dependence between random rewards and random trajectories. If different trajectories that produce identical occupancy distributions can yield different return distributions (because the reward at a given state is stochastic), then the factorization is not one-to-one. Deterministic rewards eliminate this degeneracy. The extension to finite-horizon or cycle-less environments is discussed in Appendix C.
Contrast with the standard successor measure factorization: the standard identity (Equation 2) is $V^\pi_r(x) = (1-\gamma)^{-1} \mathbb{E}_{X' \sim \Psi^\pi(\cdot \mid x)}[r(X')]$, which gives only the expected return. The distributional version (Equation 4) gives the full distribution because it replaces the deterministic expectation over $\Psi^\pi$ with the random draw from $\mathcal{M}^\pi$. In operational terms:
- Standard SR: take the successor measure, compute the expected reward under it → you get a single number (the value).
- Distributional SM: sample an occupancy distribution from
$\mathcal{M}^\pi(x)$, compute the expected reward under that sampled occupancy → you get a random draw from the return distribution. Repeat many times to build an empirical return distribution.
The Distributional Bellman Equation for Random Occupancy Measures
Having defined the target object, the paper establishes a recursive relationship that enables temporal-difference learning. Proposition 3.5 gives the distributional Bellman equation for $M^\pi$:
where $X' \sim p^\pi(\cdot \mid x)$ is a random next state drawn from the environment's transition kernel under policy $\pi$, and $S \subseteq \mathcal{X}$ is any measurable set.
What it computes: the random occupancy measure starting at $x$ has the same distribution as a mixture: with probability $(1-\gamma)$ (immediate occupancy at the current state, represented by the Dirac $\delta_x$), and with probability $\gamma$ (the discounted future occupancy, represented by the random occupancy measure starting from a randomly sampled next state $X'$). The equality is in distribution (denoted by $\stackrel{\mathcal{L}}{=}$), meaning both sides are random variables with the same probability law.
Why this form: this is a direct distributional analog of the standard Bellman equation for the successor measure. It decomposes the occupancy distribution into "now" (the Dirac at the current state, weighted by $1-\gamma$) and "later" (the occupancy from the next state onward, weighted by $\gamma$). The recursion bottoms out because $\gamma < 1$, so the influence of far-future states decays geometrically. Operationally, this equation says: if we can sample from $M^\pi(\cdot \mid X')$ for a next state $X'$, we can construct a sample from $M^\pi(\cdot \mid x)$ by taking $(1-\gamma)\delta_x$ and adding $\gamma$ times the future sample—exactly the generative procedure used to construct bootstrap targets during training.
Proof sketch: starting from Definition 3.1, separate the $k=0$ term from the sum:
Factor $\gamma$ from the remaining sum and shift the index:
By the Markov property, $(X_1, X_2, ...)$ given $X_0 = x$ has the same distribution as $(X_0, X_1, ...)$ given $X_0 = X'$ where $X' \sim p^\pi(\cdot \mid x)$. Therefore the sum on the right has the same distribution as $M^\pi(\cdot \mid X')$, yielding the equality in distribution.
From this, the paper derives a recursive expression for the distributional SM $\mathcal{M}^\pi$ itself (Equation 6):
where $b_{x,\gamma}: \mathcal{P}(\mathcal{X}) \to \mathcal{P}(\mathcal{X})$ is defined by $b_{x,\gamma}(\mu) = (1-\gamma)\delta_x + \gamma\mu$ (the function that takes an occupancy distribution $\mu$ and prepends the current state with the geometric weighting), and $f_\sharp \nu = \nu \circ f^{-1}$ is the pushforward of measure $\nu$ through function $f$ (if $Z \sim \nu$, then $f(Z) \sim f_\sharp\nu$).
What it computes: the distributional SM at state $x$ is the expected pushforward of the distributional SM at a random next state $X'$ through the mixing function $b_{x,\gamma}$. In operational terms: to get the distribution over occupancy measures starting at $x$, sample a next state $X'$ from the environment, sample an occupancy distribution $\mu$ from $\mathcal{M}^\pi(X')$ (representing where the agent goes after step 1), apply the mixing $b_{x,\gamma}$ to $\mu$ (which adds the current state to the occupancy with weight $1-\gamma$ and downweights $\mu$ by $\gamma$), and then take the expectation over $X'$.
Why this form: it defines a fixed-point equation that $\mathcal{M}^\pi$ must satisfy. This motivates the distributional Bellman operator $\mathcal{T}^\pi$ (Equation 7):
where $\mathcal{T}^\pi: \mathcal{P}(\mathcal{P}(\mathcal{X}))^\mathcal{X} \to \mathcal{P}(\mathcal{P}(\mathcal{X}))^\mathcal{X}$ maps a candidate distributional SM $\mathcal{M}$ to an updated version, and $\mathcal{P}(\mathcal{P}(\mathcal{X}))^\mathcal{X}$ is the space of functions from states to distributions over state distributions.
What it computes: given any candidate $\mathcal{M}$, the operator produces a new candidate by, at each state $x$, computing the expected pushforward of $\mathcal{M}(X')$ for $X' \sim p^\pi(\cdot \mid x)$. The true $\mathcal{M}^\pi$ is a fixed point: $\mathcal{T}^\pi \mathcal{M}^\pi = \mathcal{M}^\pi$.
Why this form: it sets up a dynamic programming approach: start with an arbitrary initial $\mathcal{M}_0$, repeatedly apply $\mathcal{T}^\pi$, and hope the iterates converge to $\mathcal{M}^\pi$. The paper proves this works.
Proposition 3.6 establishes contractivity of $\mathcal{T}^\pi$:
where $d$ is a metric on $\mathcal{X}$ making it a Polish space (complete, separable metric space), $w_d$ is the Wasserstein distance on $\mathcal{P}(\mathcal{X})$ with base distance $d$ (measuring distance between two state distributions), $W$ is the Wasserstein distance on $\mathcal{P}(\mathcal{P}(\mathcal{X}))$ with base distance $w_d$ (measuring distance between two distributions over state distributions), and $\overline{W}(\mathcal{M}_1, \mathcal{M}_2) = \sup_{x \in \mathcal{X}} W(\mathcal{M}_1(x), \mathcal{M}_2(x))$ is the supremum over states of the $W$-distance between the corresponding model distributions.
What it computes: this is a contraction factor of $\gamma < 1$ in the supremum-Wasserstein metric. When we apply $\mathcal{T}^\pi$ to two different distributional SMs, the distance between them shrinks by at least a factor of $\gamma$.
Why this form: contraction is the key property that guarantees convergence of iterative dynamic programming (via the Banach fixed-point theorem). The proof uses a novel coupling technique in the doubly-infinite-dimensional space $\mathcal{P}(\mathcal{P}(\mathcal{X}))$. The essential insight is: given an $\epsilon$-optimal coupling $\Gamma_{1,x'}$ between $\mathcal{M}_1(x')$ and $\mathcal{M}_2(x')$ with respect to $W$, we can construct a coupling between $(\mathcal{T}^\pi\mathcal{M}_1)(x)$ and $(\mathcal{T}^\pi\mathcal{M}_2)(x)$ by: (1) mixing over next states $x' \sim p^\pi(\cdot \mid x)$, (2) pushing the coupling $\Gamma_{1,x'}$ forward through $(b_{x,\gamma}, b_{x,\gamma})$. The cost of this coupling is bounded by $\gamma$ times the cost of the original coupling plus $\epsilon$, because $b_{x,\gamma}$ contracts Wasserstein distances by $\gamma$ (intuitively: mixing in the current state reduces the distinguishability of two occupancy distributions by downweighting the future component).
Corollary 3.7 establishes convergent dynamic programming: if $(\mathcal{X}, d)$ is compact (so distances are bounded), then for any initial $\mathcal{M}_0$, the iterates $\mathcal{M}_{k+1} = \mathcal{T}^\pi \mathcal{M}_k$ converge to $\mathcal{M}^\pi$ in $\overline{W}$.
What this means practically: we can, in principle, compute $\mathcal{M}^\pi$ by starting from any guess and repeatedly applying the Bellman operator. This provides theoretical justification for temporal-difference learning: the TD update approximately applies $\mathcal{T}^\pi$ using sampled transitions, and the contraction property ensures that errors from function approximation or sampling will be attenuated rather than amplified.
The δ-Model: Equally-Weighted Particle Representation
The theoretical framework operates in $\mathcal{P}(\mathcal{P}(\mathcal{X}))$, an infinite-dimensional space that cannot be represented exactly in a computer. The paper proposes the equally-weighted particle (EWP) representation:
where $m$ is the number of particles (model atoms), $\theta_i(x) \in \mathcal{P}(\mathcal{X})$ is the $i$-th generative model of state conditioned on starting state $x$ (each $\theta_i(x)$ is itself a probability distribution over $\mathcal{X}$), and $\delta_{\theta_i(x)}$ is a Dirac mass at distribution $\theta_i(x)$ in the space $\mathcal{P}(\mathcal{X})$.
What it represents: the continuous distribution $\mathcal{M}^\pi(x)$ over occupancy distributions is approximated by a discrete uniform distribution over $m$ representative occupancy distributions. Each $\theta_i(x)$ is one possible "future" that the agent might experience—one specific discounted visitation pattern. The collection of $m$ such futures, each equally likely, forms a Monte Carlo approximation to the true distribution of futures.
Why this form: this is a direct analog of the quantile representation used in standard distributional RL (Dabney et al., 2018b; Nguyen-Tang et al., 2021), where a return distribution is represented as $\frac{1}{m} \sum_{i=1}^m \delta_{z_i}$ for scalar quantile locations $z_i$. The difference here is that each atom is not a scalar but a full generative model of state—a distribution over $\mathcal{X}$. This is necessary because $\mathcal{M}^\pi(x)$ lives in $\mathcal{P}(\mathcal{P}(\mathcal{X}))$, not $\mathcal{P}(\mathbb{R})$.
Why generative models for the atoms: each $\theta_i(x)$ must be a distribution over $\mathcal{X}$ that we can sample from (to compute expected rewards) and that we can compare to other distributions (to define distances). Generative models—specifically, neural networks that take noise variables and output state samples—satisfy both requirements. The paper draws a direct connection to $\gamma$-models (Janner et al., 2020), which learn a single generative model of the successor measure. The δ-model generalizes this to an ensemble, but with the crucial difference that the ensemble members are coupled through the model-level MMD loss rather than trained independently:
"We propose to represent the atoms as generative models, in the spirit of γ-models (Janner et al., 2020). In practice, the generative models can be implemented with function approximators that take as input noise variables similar to the generator of a generative adversarial network (GAN; Goodfellow et al., 2014). We refer to such an EWP model as a δ-model; Figure 2 illustrates its components."
Terminology established by the paper (Section 4.1):
- Model distribution
$\hat{\mathcal{M}}(x)$: the distribution over generative models (the ensemble), living in$\mathcal{P}(\mathcal{P}(\mathcal{X}))$. - State distributions or model atoms
$\theta_i(x)$: the individual generative models, each living in$\mathcal{P}(\mathcal{X})$. - Model sample: a generative model
$\theta \sim \hat{\mathcal{M}}(x)$drawn from the ensemble (i.e., selecting one of the$m$atoms uniformly). - State sample: a state
$X' \sim \theta$drawn from a generative model.
Operational interpretation of the two levels:
- Level 1 (between model atoms): the distribution
$\hat{\mathcal{M}}(x)$captures which futures are possible. If the policy can go left or right in a maze, two atoms might represent "mostly left-future occupancy" and "mostly right-future occupancy." - Level 2 (within each atom): each
$\theta_i(x)$captures what states are visited within that future. The "mostly left-future" atom puts high probability on states in the left corridor and low probability on states in the right corridor.
The Two-Level MMD Loss Construction
To train the δ-model, the paper needs a loss function that compares the current model distribution $\hat{\mathcal{M}}(x)$ to a target distribution constructed from the distributional Bellman equation. This requires defining a distance between elements of $\mathcal{P}(\mathcal{P}(\mathcal{X}))$—distributions over distributions. The paper constructs this in three layers.
Layer 1: State kernel $\kappa$ for comparing individual state distributions.
The maximum mean discrepancy (MMD; Gretton et al., 2012) between two probability distributions $p, q$ over a set $\mathcal{Y}$ with kernel $\kappa: \mathcal{Y} \times \mathcal{Y} \to \mathbb{R}$ is:
where $(X, X') \sim p \otimes p$ are two independent samples from $p$, and $(Y, Y') \sim q \otimes q$ are two independent samples from $q$.
What it computes: the MMD measures the squared distance between the mean embeddings of $p$ and $q$ in the reproducing kernel Hilbert space (RKHS) induced by $\kappa$. The first term $\mathbb{E}[\kappa(X, X')]$ is the self-similarity of $p$ (average kernel value between two independent draws from $p$); the second term $\mathbb{E}[\kappa(Y, Y')]$ is the self-similarity of $q$; the cross term $-2\mathbb{E}[\kappa(X, Y)]$ is the similarity between samples from $p$ and $q$. If $p$ and $q$ are identical, these three terms cancel to zero. If they differ, the self-similarities exceed the cross-similarity, yielding a positive distance.
Why this form: MMD is a metric when $\kappa$ is characteristic (injective mean embedding), providing a proper distance between probability distributions. Unlike the Wasserstein distance, MMD can be estimated efficiently from samples without solving an optimal transport problem. For the δ-model, where we represent each $\theta_i(x)$ as a generative model, we can estimate the MMD between any two model atoms by drawing $n_1$ samples from one generator and $n_2$ samples from the other.
The sample-based estimator used in the paper (Equation 10), with $X_1, ..., X_{n_1} \overset{\text{i.i.d.}}{\sim} \theta$ and $Y_1, ..., Y_{n_2} \overset{\text{i.i.d.}}{\sim} \theta'$, is:
where $\binom{n}{2} = n(n-1)/2$ is the number of unordered pairs, the first term $\frac{\sum_{i<j} \kappa(X_i, X_j)}{\binom{n_1}{2}}$ is the U-statistic estimator of the self-similarity of $\theta$ (averaging kernel values over all distinct pairs of samples from $\theta$), the second term is the analogous estimator for $\theta'$, and the third term is the average kernel value across all cross-pairs between the two sets of samples.
What it computes: an unbiased estimate of the squared MMD between the two generative models, using only samples from each. No density evaluation or integration is required—only the ability to sample and compute kernel evaluations between samples.
Why this form: the U-statistic form (using $\binom{n}{2}$ in the denominator rather than $n^2$) removes the diagonal terms $\kappa(X_i, X_i)$ which would bias the estimate upward (since a sample is perfectly similar to itself). The cross-term uses all $n_1 \times n_2$ pairs because samples from different distributions are always independent.
Layer 2: Model kernel $k_\kappa$ for comparing distributions over state distributions.
The MMD defined above gives a distance $\text{MMD}_\kappa(\theta, \theta')$ between two state distributions. To define a kernel on $\mathcal{P}(\mathcal{X})$ itself—which will allow us to compute an MMD between model distributions in $\mathcal{P}(\mathcal{P}(\mathcal{X}))$—the paper follows Christmann & Steinwart (2010) and Szabo et al. (2015) by composing a radial basis function with the state MMD:
where $\sigma > 0$ is a bandwidth parameter controlling the scale of similarity, and $\rho(y) = (1 + y^2)^{-1/2}$ is the inverse multiquadric radial basis function.
What it computes: for any two state distributions $\theta, \theta'$, this kernel outputs a scalar in $(0, 1]$ indicating their similarity. When $\theta$ and $\theta'$ are identical, $\text{MMD}_\kappa(\theta, \theta') = 0$, so $\rho(0) = 1$. As the MMD grows, $\rho$ decays toward zero, with the bandwidth $\sigma$ controlling how quickly similarity drops off.
Why the inverse multiquadric: Szabo et al. (2015, Table 1) shows that for this choice of $\rho$, the kernel $k_\kappa$ is characteristic on $\mathcal{P}(\mathcal{X})$. This means the resulting MMD over $\mathcal{P}(\mathcal{P}(\mathcal{X}))$ is a proper metric—it is zero if and only if the two model distributions are identical. This is essential for the loss to have a unique minimum at the true $\mathcal{M}^\pi$.
Layer 3: The DSM MMD loss (Equation 12).
With the model kernel $k_\kappa$ in hand, the paper defines the loss as the MMD between the current δ-model $\hat{\mathcal{M}}(x)$ and the Bellman target $\tilde{\mathcal{M}}(x)$:
where $\hat{\mathcal{M}}(x) = \frac{1}{m}\sum_{i=1}^m \delta_{\theta_i(x)}$ is the current δ-model at state $x$, and $\tilde{\mathcal{M}}(x)$ is the target distribution defined by:
What the target represents: this is the empirical Bellman target. For a sampled transition $(x, x')$, each model atom $\theta_i(x')$ (representing a future occupancy distribution from $x'$) is transformed via $b_{x,\gamma}$ into $(1-\gamma)\delta_x + \gamma\theta_i(x')$—the occupancy distribution that starts at $x$, then follows the trajectory that $\theta_i(x')$ represents from $x'$ onward. The mixture over $i$ with equal weights and the expectation over $x'$ approximates the right-hand side of the distributional Bellman equation.
Why this form for the loss: the MMD under the model kernel $k_\kappa$ measures the distance between two distributions over $\mathcal{P}(\mathcal{X})$ in a way that is: (a) a proper metric (since $k_\kappa$ is characteristic); (b) computable from samples (since we can sample model atoms from each distribution); (c) amenable to gradient-based optimization (since MMD is a smooth function of its inputs when the kernels are differentiable). Minimizing this loss drives $\hat{\mathcal{M}}(x)$ toward satisfying the distributional Bellman equation.
Why MMD rather than other distribution distances: Wasserstein distance would require solving an optimal transport problem over $\mathcal{P}(\mathcal{X})$ at each training step—computationally infeasible. Kullback-Leibler divergence would require density estimation over $\mathcal{P}(\mathcal{X})$, which is ill-defined for the particle representation. MMD avoids both issues: it can be estimated from samples with a closed-form expression, and it is well-defined for empirical distributions (sums of Diracs).
Expanded loss for the δ-model representation. Substituting the particle representation $\hat{\mathcal{M}}(x) = \frac{1}{m}\sum_i \delta_{\theta_i(x)}$ and the single-sample target (using one sampled $x'$ in place of the expectation), the MMD expands as:
where $\theta_i(x)$ are the current model atoms, $\bar{\theta}_j(x) = (1-\gamma)\delta_x + \gamma\theta_j(x')$ are the target atoms after applying the Bellman transformation, and we have dropped the constant term $\frac{1}{m^2}\sum_{i,j} k_\kappa(\bar{\theta}_i(x), \bar{\theta}_j(x))$ because it does not depend on the current model parameters (it only depends on the target atoms, which are computed using target network parameters—a fixed target, as in DQN).
What this expands to operationally: the loss is the difference between the average within-source similarity (how similar the current atoms are to each other) and twice the average cross-similarity (how similar each current atom is to each target atom). The omitted target self-similarity term is constant with respect to the current parameters. Minimizing this loss pushes each $\theta_i(x)$ to be similar to some target atom $\bar{\theta}_j(x)$ while maintaining diversity among the $\theta_i(x)$ (since excessive similarity among current atoms would make the first term large, which is penalized).
Computing the kernel evaluations in practice. Each $k_\kappa(\theta, \theta')$ evaluation requires computing $\text{MMD}_\kappa(\theta, \theta')$, which is estimated by drawing $s$ state samples from each generative model. The paper specifies this in Algorithm 1, with a default of $s = 32$ state samples per model atom. For $m$ model atoms, the total number of kernel evaluations per state $x$ scales as $\mathcal{O}(m^2 s^2)$.
Sampling from the target atoms. A subtle but important practical detail: to sample from $\bar{\theta}_j(x) = (1-\gamma)\delta_x + \gamma\theta_j(x')$, we cannot directly draw from a mixture of a Dirac and a generative model. The paper's solution (described in Section 4.2 and made explicit in Algorithm 1 for the $n$-step case):
"we can sample from distributions of the form
$(1-\gamma)\delta_x + \gamma\theta_i(x')$by first sampling$Y \sim \text{Bernoulli}(1-\gamma)$, returning$x$if$Y = 1$, and otherwise returning an independent sample from$\theta_i(x')$."
What this does: with probability $1-\gamma$, the sample is the current state $x$ (the "now" component). With probability $\gamma$, the sample is drawn from the generative model $\theta_i(x')$ (the "future" component). This procedure exactly simulates the mixture distribution and is used both during training (to build targets) and during inference (to compute expected rewards).
n-Step Bootstrapping
The one-step Bellman target described above has a bootstrap fraction of $\gamma$ (the mass placed on $\theta_i(x')$). For typical discount factors like $\gamma = 0.95$, the bootstrap component accounts for 95% of the target distribution's mass. When the δ-model is early in training and $\theta_i(x')$ produces poor samples, the target is dominated by a low-quality signal, leading to instability.
The paper addresses this with n-step bootstrapping, extending Equation 5 by unrolling the recursion:
where $X_0 = x$, $X_{k+1} \sim p^\pi(\cdot \mid X_k)$ for $k = 0, ..., n-1$, and $n$ is the number of unrolled steps.
What it computes: the random occupancy measure is equal in distribution to a mixture of: (1) $n$ Dirac masses at the first $n$ states of the trajectory, weighted by $(1-\gamma)\gamma^i$ for $i = 0, ..., n-1$ (the "ground-truth" component based on actual observed states), and (2) the bootstrapped future occupancy from state $X_n$, weighted by $\gamma^n$ (the "learned" component).
Why this helps: for $\gamma = 0.95$ and $n = 5$, the bootstrap weight is $\gamma^5 \approx 0.77$ (reduced from 0.95). The remaining 23% of the target is grounded in actual observed trajectory states, providing a stronger learning signal. The paper reports:
"We found that training stability tends to improve substantially when bootstrap samples account for roughly 80% of the samples in the procedure above."
Appendix E includes an ablation (Figure 9, left) on Pendulum showing that Wasserstein distance to the true return distribution improves as $n$ increases from 1 to about 5, after which benefits level off. The default hyperparameter used is $n = 5$ (Table 1).
Sampling from the n-step target (Algorithm 1):
"we can sample from this distribution by first sampling
$Y$from a$\text{Geometric}(1-\gamma)$distribution, returning$x_k$if$Y = k < n$, and returning a sample from$\theta_i(x_n)$otherwise."
What this does: the geometric distribution with parameter $1-\gamma$ places probability $(1-\gamma)\gamma^k$ on outcome $k$—exactly matching the weights in the n-step expansion. With probability $\sum_{k=0}^{n-1} (1-\gamma)\gamma^k = 1 - \gamma^n$, we return an observed trajectory state $x_k$. With probability $\gamma^n$, we return a sample from the bootstrapped model at $x_n$.
This technique is not specific to the distributional SM; the paper notes that it should be broadly useful for training geometric horizon models like $\gamma$-models, where long horizons were reported as a major challenge (Janner et al., 2020; Thakoor et al., 2022).
Adaptive and Adversarial Kernels
The non-stationarity of the learning process—the model atoms $\theta_i(x)$ are continually evolving—poses a challenge for kernel selection. A fixed kernel with a fixed bandwidth may be a poor similarity measure early in training (when the generators produce crude samples) versus late in training (when generators are more refined).
Adversarial state kernel. The paper employs the MMD-GAN approach (Li et al., 2017; Binkowski et al., 2018), where the state kernel is parameterized by a learned feature map:
where $f: \mathcal{X} \to \mathcal{Z}$ is a learned embedding function (mapping raw states to a feature space), and $\kappa$ is a fixed base kernel on the feature space (in practice, a mixture of rational quadratic kernels with bandwidths $A = \{0.2, 0.5, 1.0, 2.0, 5.0\}$).
What this does: instead of computing kernel similarities directly on raw states, states are first transformed through $f$, and the kernel is evaluated in the transformed space. The function $f$ is trained adversarially to maximize the model MMD (making it better at distinguishing different model distributions), while the generators are trained to minimize the same MMD. This is the GAN-like min-max dynamic: the discriminator (kernel) tries to detect differences between the current δ-model and the target; the generator tries to make them indistinguishable.
Why injectivity matters: Li et al. (2017) shows that if $f$ is injective (one-to-one) and $\kappa$ is characteristic on $\mathcal{Z}$, then $\kappa_f$ is characteristic on $\mathcal{X}$. A non-injective $f$ could collapse distinct states to the same feature vector, losing the ability to distinguish different occupancy distributions. The standard MMD-GAN approach encourages injectivity by training $f$ as the encoder of an autoencoder, where the reconstruction loss penalizes information loss.
The paper's innovation: iResNet for guaranteed injectivity. Rather than relying on the autoencoder's reconstruction loss as a soft constraint toward injectivity, the paper uses an invertible residual network (iResNet; Behrmann et al., 2019):
"Unlike an autoencoder, this enforces injectivity, and to our knowledge, no other work has employed invertible neural networks for modeling an adversarial kernel."
What this is: an iResNet is a neural network architecture where each layer is designed to be invertible—the forward mapping $x \mapsto f(x)$ has a well-defined inverse $f^{-1}$. This guarantees injectivity by construction, since if $f(x) = f(y)$, then applying the inverse gives $x = y$. The architecture uses constrained residual blocks where the Lipschitz constant of the residual branch is strictly less than 1, ensuring the overall transformation is a homeomorphism.
Why this matters for the DSM: guaranteed injectivity means the adversarial kernel remains characteristic throughout training, regardless of how aggressively the discriminator is optimized. With an autoencoder-based approach, if the discriminator "wins" too hard, the encoder might learn to collapse states, degrading the kernel's ability to measure meaningful distances. The iResNet prevents this failure mode entirely.
Adaptive model kernel bandwidth. Even with an adaptive state kernel, the model kernel bandwidth $\sigma$ (which scales the MMD values before passing through $\rho$) needs adaptation. The paper uses the median heuristic (Takeuchi et al., 2006; Gretton et al., 2012):
"Prior to computing the model MMD, we choose
$\sigma^2$to be the median of the pairwise$\text{MMD}^2_\kappa$between the model atoms of$\hat{\mathcal{M}}(x)$and those of the bootstrap target$\tilde{\mathcal{M}}(x)$."
What this computes: take all pairwise state-MMD values between the $m$ source atoms $\theta_i(x)$ and the $m$ target atoms $\bar{\theta}_j(x)$, plus the within-source and within-target pairs (a total of $\mathcal{O}(m^2)$ distances). Set $\sigma^2$ to the median of these squared MMD values.
Why this form: the median heuristic adapts the kernel's sensitivity to the current scale of variation in the data. If the atoms are very similar (small MMD values), $\sigma$ will be small, making the model kernel more sharply peaked and more sensitive to small differences. If the atoms are very different (large MMD values), $\sigma$ will be large, preventing the kernel from saturating near zero. This automatic scaling is crucial because the typical distance between occupancy distributions changes dramatically over the course of training.
Algorithm 1 implementation detail. The full algorithmic procedure in Appendix A shows that the MMDs are computed in three blocks: $d^s_{i,i'}$ (source within-source MMDs), $d^t_{i,i'}$ (target within-target MMDs), and $d^{st}_{i,i'}$ (source-target cross MMDs). The median $\sigma^2$ is computed over the concatenation of all three sets of distances. The model MMD loss is then:
What this computes: this is exactly the expansion of $\text{MMD}^2_{k_\kappa}(\hat{\mathcal{M}}(x), \tilde{\mathcal{M}}(x))$ with the particle representation, where each $\text{MMD}_\kappa(\theta, \theta')$ value is passed through $\rho(\cdot/\sigma)$ after the adaptive bandwidth normalization.
Ablation evidence (Appendix E, Figure 9 right). The paper reports that removing the adversarial kernel (using a fixed mixture of rational quadratic kernels directly on raw states) approximately doubles the Wasserstein distance between the predicted and true return distributions. This confirms that adaptive kernels are not just a minor optimization but a critical component for achieving accurate return distribution estimates.
Summary of Design Choices and Their Justifications
-
Distributions over distributions vs. scalar distributions: placing atoms in
$\mathcal{P}(\mathcal{X})$rather than a Euclidean space is necessary because the random occupancy measure is a distribution over states, not a vector. Attempting to embed occupancy distributions into a finite-dimensional space (e.g., via sufficient statistics) would lose information about the joint visitation patterns that are essential for correct return distribution estimation (as the prior work on distributional successor features fails to capture joint distributions). -
Equally-weighted particles with generative model atoms: the EWP representation inherits from the quantile distributional RL literature, providing a non-parametric, sample-based approximation whose resolution scales with
$m$. Using generative models (rather than, e.g., Gaussian mixtures with fixed parametric form) maintains flexibility to represent arbitrary occupancy distributions, including multi-modal ones arising from stochastic policies in branching environments (Figure 4a demonstrates that the DSM captures distinct behavioral modes that an ensemble of independently-trained$\gamma$-models collapses to the mean). -
Two-level MMD rather than a single-level loss: a loss defined directly on
$\mathcal{P}(\mathcal{X})$(e.g., MMD between the average occupancy distribution and the average target occupancy) would only match first moments—it would drive the mean of the δ-model to match the successor measure, but would not constrain higher moments or the diversity of atoms. The two-level MMD compares the full model distributions in$\mathcal{P}(\mathcal{P}(\mathcal{X}))$, ensuring that the distribution over futures (not just the average future) is accurately captured. -
Inverse multiquadric for the model kernel: chosen for its characteristic property on
$\mathcal{P}(\mathcal{X})$(Szabo et al., 2015). Alternative choices like the Gaussian RBF$\rho(y) = \exp(-y^2/2)$would also be characteristic, but the inverse multiquadric has heavier tails (decays polynomially rather than exponentially), which may provide more stable gradients when the MMD values are large. -
Median heuristic for bandwidth: adapts to the non-stationary scale of the occupancy distributions during training without introducing additional learned parameters. An alternative—learning
$\sigma$as a parameter—would add optimization complexity and potential for degenerate solutions (e.g.,$\sigma \to 0$making the kernel a delta function, which trivially minimizes the MMD but loses all metric structure). -
iResNet for adversarial kernel injectivity: guarantees the state kernel remains characteristic throughout adversarial training without relying on auxiliary reconstruction losses that may conflict with the discriminator objective. This is a novel application of invertible networks in the context of kernel learning.
-
$n$-step bootstrapping with$n$chosen so bootstrap mass$\approx$80%: balances the bias-variance tradeoff. Too small$n$(high bootstrap mass) yields unstable learning due to noisy bootstrap targets early in training. Too large$n$requires storing long trajectory segments in memory and reduces the effective number of independent updates per trajectory. The 80% heuristic is empirically derived but matches the intuition that the targets should be mostly grounded in learned structure (since that's what the model needs to improve) but with sufficient real-data signal to anchor the learning. -
Fixed target parameters with periodic updates (Algorithm 1,
$\lambda = 0.01$): the target model atoms$\bar{\zeta}_i$are updated via Polyak averaging$\bar{\zeta}_i \leftarrow (1-\lambda)\bar{\zeta}_i + \lambda\zeta_i$, a standard technique from deep RL (Mnih et al., 2015) that stabilizes TD learning by making the target distribution change slowly relative to the online model. -
Default hyperparameters (Table 1):
$m = 51$model atoms (Appendix D.3.2 for Pendulum),$s = 32$state samples per atom,$n = 5$step unrolling,$\gamma = 0.95$, generator and discriminator as 3-layer MLPs with 256 units and ReLU activations, discriminator feature dimensionality of 8 output features (the embedding dimension for the adversarial kernel), Adam optimizer with$\beta_1 = 0.9, \beta_2 = 0.999$and learning rate$6.25 \times 10^{-5}$for both generator and discriminator,$\lambda = 0.01$for target network updates, and$3 \times 10^6$total gradient updates for Pendulum experiments.
4. Key Insights and Innovations
Innovation 1: The Random Occupancy Measure as a Distributional Analog to the Successor Representation—a Conceptual Bridge Between Two Previously Disconnected Frameworks
The paper's most fundamental intellectual move is not algorithmic but definitional: identifying which mathematical object should serve as the distributional counterpart to the successor representation. The standard successor measure $\Psi^\pi$ factorizes expected returns as a linear operator applied to a reward function. Prior work attempting to bring distributional thinking to this space—such as distributional successor features (Gimelfarb et al., 2021; Carvalho et al., 2023)—essentially learned distributions over the output of this operator (distributions over scalar features), but did so by modeling marginal distributions independently, discarding the joint structure that makes the factorization meaningful.
The innovation here is to step back and ask: what is the input to the linear operator whose distribution we should be modeling? The answer—the random occupancy measure $M^\pi$—is a conceptual shift that the paper makes look natural in retrospect but which required recognizing that the successor measure itself is an expectation of a random object (Equation 3). By "removing the expectation" from the definition of $\Psi^\pi$, the paper reveals a richer object that lives in $\mathcal{P}(\mathcal{P}(\mathcal{X}))$ rather than $\mathcal{P}(\mathcal{X})$. This is not an incremental extension of the successor representation; it is a fundamentally different type of object—a distribution over probability measures rather than a single probability measure.
The significance of this framing extends beyond the specific algorithm in the paper. It provides a template for distributional analogs of other linear operators in reinforcement learning: identify the random variable whose expectation is the operator of interest, and model its distribution. This same pattern could apply to successor features, to goal-conditioned value functions, or to any setting where a linear factorization between dynamics and reward is desirable. The paper does not explore these extensions, but the conceptual framework enables them.
The comparison to an ensemble of $\gamma$-models (Section 6) crystallizes why this distinction matters. An ensemble of independently-trained successor measure models captures epistemic uncertainty—our ignorance about the true $\Psi^\pi$ given finite data. As data increases, all ensemble members converge to the same distribution, and the ensemble collapses. The distributional SM captures aleatoric uncertainty—the inherent randomness in which trajectory the agent actually experiences, even given perfect knowledge of the environment. This randomness does not vanish with more data; it is a property of the stochastic policy and environment. The experimental demonstration (Figure 4a) that the DSM recovers distinct behavioral modes (different corridors in the Windy Gridworld) while the $\gamma$-model ensemble collapses to the mean is not merely a performance improvement—it is evidence that the DSM is measuring something the ensemble fundamentally cannot, regardless of ensemble size or data quantity.
This is a fundamental conceptual advance rather than an incremental improvement. It changes what it means to do distributional RL with successor-type representations: from "learn distributions over scalar features" to "learn distributions over occupancy measures."
Innovation 2: A Distributional Bellman Equation Operating on $\mathcal{P}(\mathcal{P}(\mathcal{X}))$—Extending the Scope of Distributional Dynamic Programming to Doubly-Infinite-Dimensional Spaces
Distributional Bellman equations are well-established in standard distributional RL, where the return distribution $\eta^\pi(x) \in \mathcal{P}(\mathbb{R})$ satisfies a recursive relationship involving a pushforward through a reward-plus-discounting operator (Bellemare et al., 2017a; Morimura et al., 2010). These operate on distributions over the real line—objects that, while infinite-dimensional in principle, are parameterized in practice by finite sets of quantiles or categorical bins.
The distributional SM Bellman equation (Proposition 3.5, Equation 5) operates on a qualitatively different space: $\mathcal{P}(\mathcal{P}(\mathcal{X}))$—distributions over state distributions. If $\mathcal{X}$ is continuous, this is a doubly infinite-dimensional space. Prior distributional RL theory provided no guidance on whether Bellman operators on such spaces would be well-behaved—whether they contract, whether dynamic programming converges, or whether the fixed point is unique.
The paper's theoretical contribution in Propositions 3.6 and Corollary 3.7 is to establish that the distributional Bellman operator $\mathcal{T}^\pi$ is a $\gamma$-contraction in the supremum-Wasserstein metric $\overline{W}$ over $\mathcal{P}(\mathcal{P}(\mathcal{X}))^\mathcal{X}$, and that iterative application converges to the unique fixed point $\mathcal{M}^\pi$. This is not a trivial extension of existing proofs. The standard distributional RL contraction proof (Bellemare et al., 2017a) relies on the structure of the pushforward through a reward-plus-discount function $f(r, \gamma) = r + \gamma z$, using couplings over $\mathbb{R}$. The DSM proof requires constructing couplings over $\mathcal{P}(\mathcal{X})$ and bounding the Wasserstein distance $w_d$ under the mixing operator $b_{x,\gamma}(\mu) = (1-\gamma)\delta_x + \gamma\mu$. The paper develops a novel coupling argument (Appendix B.1) showing that $b_{x,\gamma}$ contracts $w_d$ by $\gamma$, which then propagates upward through the nested Wasserstein construction.
Why this matters beyond the specific result: it establishes that distributional dynamic programming scales to complex structured spaces, not just scalar returns. The proof technique—which the authors note draws on the coupling approach of Amortila et al. (2020)—provides a template for analyzing distributional Bellman operators on spaces of distributions over general metric spaces. This opens the door to distributional analogs of other structured prediction problems in RL (multi-task distributions, goal-conditioned distributions, trajectory-level distributions) that were previously theoretically uncharted.
The contractivity result also provides theoretical justification for the temporal-difference learning approach taken in the paper. The $\gamma$-contraction property means that errors from function approximation, sampling noise, and optimization imperfections are attenuated rather than amplified by the Bellman operator—a property that is not guaranteed a priori for such a complex space, and whose absence would make learning infeasible regardless of representational capacity.
This is a theoretical advance that extends the foundations of distributional RL to a new class of objects and provides the rigorous underpinning needed for the algorithmic contributions to be meaningful.
Innovation 3: The Two-Level MMD as a Loss on $\mathcal{P}(\mathcal{P}(\mathcal{X}))$—A Practical Distance for Doubly-Structured Distributions That Avoids the Pitfalls of Competing Approaches
The problem of defining a loss function over $\mathcal{P}(\mathcal{P}(\mathcal{X}))$—distributions over probability distributions—is not one that has a standard solution in the machine learning toolkit. The paper's approach of using a two-level MMD, where a kernel on $\mathcal{P}(\mathcal{X})$ is constructed from an MMD on $\mathcal{X}$, and then an MMD on $\mathcal{P}(\mathcal{P}(\mathcal{X}))$ is computed using that kernel, is technically elegant. But its true innovation is in why it works when plausible alternatives fail.
Consider the alternatives. One could embed each occupancy distribution into a finite-dimensional summary statistic (e.g., moments, expected features) and define a loss over the resulting Euclidean space. This would be a distributional successor features approach (Gimelfarb et al., 2021), which the paper explicitly identifies as insufficient because it discards joint distributional information—two different occupancy distributions can have identical marginal feature expectations but different joint visitation patterns, leading to different return distributions for correlated reward functions.
One could use a GAN-style discriminator directly on $\mathcal{P}(\mathcal{P}(\mathcal{X}))$—a neural network that takes as input a set of state samples and classifies whether they come from the true or generated model distribution. This would require a set encoder architecture capable of handling variable-size sets of variable-size sets (each model atom produces multiple state samples, and there are multiple model atoms), and would lose the metric guarantees that MMD provides. It would also be vulnerable to mode collapse in ways that MMD, with its characteristic kernel, is more robust to.
One could use the Wasserstein distance at both levels, but computing the Wasserstein distance over $\mathcal{P}(\mathcal{X})$ requires solving an optimal transport problem at each training step, which is computationally prohibitive.
The MMD construction sidesteps all of these issues. At the state level, MMD provides a closed-form, sample-based distance between generative models without density estimation or optimal transport—only kernel evaluations are needed. At the model level, the inverse multiquadric composition $\rho(\text{MMD}_\kappa(\theta, \theta') / \sigma)$ converts the state MMD into a characteristic kernel on $\mathcal{P}(\mathcal{X})$ (following Szabo et al., 2015), ensuring the outer MMD is a proper metric on $\mathcal{P}(\mathcal{P}(\mathcal{X}))$. The resulting loss is a smooth function of generator parameters, enabling gradient-based optimization, while maintaining the metric property that the loss is zero if and only if the model distribution matches the target.
What makes this more than just a technical specification is the diagnostic insight it embodies: the paper recognizes that the right distance metric must respect the two-level structure of the problem, and that constructing it compositionally—state distance → model kernel → model distance—preserves the properties (characteristicness, differentiability, sample-based estimability) needed at each level. This compositional approach to defining distances on nested spaces of distributions is a design pattern that could apply beyond this paper to any setting involving distributions of structured objects (distributions of trajectories, distributions of policies, distributions of environments).
The experimental ablation in Appendix E (Figure 9, right) provides evidence that this choice is not arbitrary: removing the adaptive kernel (using a fixed mixture of rational quadratic kernels directly on raw states without the learned embedding) approximately doubles the Wasserstein distance to the true return distribution. This suggests that the metric structure matters substantially, not just the overall loss family.
This is a methodological contribution that provides a blueprint for defining learning objectives on spaces of distributions-over-distributions, with the compositional MMD approach serving as a practical template for future work on structured distributional learning problems.
Innovation 4: Zero-Shot Risk-Sensitive Policy Selection—A Capability That No Prior Method Could Provide, Enabled by the Distributional Factorization
While the distributional SM is presented primarily as a theoretical construct for factorizing return distributions, the paper demonstrates a downstream capability that serves as an existence proof for the framework's practical value: zero-shot risk-sensitive policy comparison across novel reward functions. This is not an incremental improvement over existing methods—it is a capability that no prior approach, whether from distributional RL, successor representations, or model-based RL, could provide.
To understand why this is distinctive, consider what each existing approach can and cannot do:
-
Standard distributional RL (Bellemare et al., 2017a; Dabney et al., 2018b): can compute return distributions for a given reward function after training, enabling risk-sensitive evaluation of a single policy on a single task. Cannot evaluate new reward functions without retraining, because the return distribution is tied to the reward encountered during learning.
-
Successor measures /
$\gamma$-models (Dayan, 1993; Janner et al., 2020): can perform zero-shot evaluation of novel reward functions by reweighting occupancy predictions. But this yields only expected returns—a single scalar per state—which can rank policies by mean performance but cannot compare them on risk-sensitive criteria like conditional value-at-risk, variance, or probability of catastrophic failure. -
Learned transition models
$\hat{p}^\pi$: in principle can sample rollouts for any reward function and compute empirical return distributions. The paper demonstrates that this fails in practice due to compounding model error over long horizons (Figure 3b, "Rollout" baseline), consistent with well-documented challenges in model-based RL (Jafferjee et al., 2020; Lambert et al., 2022). Even if it worked, the computational cost of sampling full trajectories for each new reward function would be substantial. -
Distributional successor features (Gimelfarb et al., 2021; Carvalho et al., 2023): can compute risk-sensitive criteria for novel tasks if those criteria depend only on marginal feature distributions. Cannot compute the full return distribution for an arbitrary reward function because the joint distribution over features is not modeled.
The distributional SM occupies the intersection of capabilities that no prior method reaches: it provides zero-shot transfer across reward functions (like successor measures) while preserving the full distributional information needed for risk-sensitive evaluation (like distributional RL, but without the reward dependence). Figure 4b demonstrates this concretely: for a held-out "Lopsided Checkerboard" reward function, the DSM correctly identifies that policy $\pi_1$ is superior under the mean return criterion while policy $\pi_2$ is superior under the 0.4-CVaR criterion—a ranking that no successor measure (which only sees means) and no standard distributional RL method (which cannot evaluate new rewards zero-shot) could produce. The paper highlights this unambiguously:
"We note that, to our knowledge, no other method can accomplish this feat. On the one hand, existing distributional RL algorithms could not evaluate the return distributions for held-out reward functions. On the other hand, any algorithm rooted in the SM for zero-shot evaluation can only rank policies by their mean returns, so they must fail to rank
$\pi_1$,$\pi_2$by at least one of the objectives."
The significance of this capability extends beyond the specific experimental demonstration. It suggests a deployment model where a policy is trained once with a general behavioral objective, its distributional SM is learned, and then at decision time—perhaps in response to changing environmental conditions, evolving safety constraints, or stakeholder preferences—the policy can be evaluated against multiple risk-sensitive criteria simultaneously, with the appropriate policy selected on-the-fly. This decouples policy training (which requires expensive interaction) from policy evaluation and selection (which becomes a cheap post-hoc computation), a separation that is particularly valuable in safety-critical settings where the risk criteria may not be fully known at training time or may change as new hazards are discovered.
The paper does not claim to solve the full problem of risk-sensitive policy optimization under novel reward functions—it only evaluates pre-existing policies. But the capability it demonstrates is the necessary first step: if you cannot evaluate, you cannot optimize. The factorization provided by the distributional SM makes risk-sensitive policy search computationally feasible by reducing it to repeated applications of Equation 4 rather than repeated environment interaction.
This is a capability-level contribution enabled by the theoretical framework. It transforms the distributional SM from a mathematical curiosity into a tool with a clear, previously-unattainable practical use case.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The paper evaluates on two continuous control domains: Windy Gridworld, a stochastic 2D navigation environment where a pointmass moves on a continuous grid subject to random wind forces pushing it toward the corners, and Pendulum (Atkeson & Schaal, 1997), a classic continuous control task requiring a policy to swing up and balance a pendulum. Pendulum serves as the primary quantitative benchmark; Windy Gridworld is used for qualitative visualizations of model atoms and the risk-sensitive policy selection demonstration. No standard train/test split is reported—the environments are used as generated via simulation.
-
Base model(s). The δ-model uses a generator network (3-layer MLP, 256 units, ReLU) that takes as input a source state
$x$and a noise vector$\omega \in \mathbb{R}^8 \sim \mathcal{N}(0, I)$, and outputs a state sample from the corresponding occupancy distribution. The adversarial kernel uses an iResNet (invertible residual network; Behrmann et al., 2019) with 2 layers × 2 blocks, 256 units, ReLU activations, projecting to 8 output features. Both networks are optimized with Adam ($\beta_1 = 0.9$,$\beta_2 = 0.999$, learning rate$6.25 \times 10^{-5}$). The δ-model uses$m = 4$model atoms for Windy Gridworld and$m = 51$for Pendulum. -
Metrics. The primary evaluation metric is the Wasserstein distance between the return distribution predicted by the DSM (via Equation 4) and the empirical Monte Carlo (MC) return distribution estimated from environment rollouts. For the risk-sensitive policy selection experiment (Figure 4b), the paper evaluates mean return and conditional value-at-risk at level 0.4 (0.4-CVaR; Rockafellar & Uryasev, 2002), comparing DSM estimates to MC ground truth. The Cramér distance (Székely & Rizzo, 2013; Bellemare et al., 2017b) is also reported in the return distribution quality comparison (Figure 3b).
-
Baselines. Four baselines are compared:
- γ-Model Ensemble: an ensemble of
$m$independently-trained$\gamma$-models (Janner et al., 2020), where each ensemble member is structurally equivalent to one δ-model atom but trained with a sum of individual state-MMD losses rather than the coupled model-MMD loss. The ensemble size$m$matches the number of atoms in the comparable δ-model. - Rollout (learned transition kernel): a learned model of the one-step transition kernel
$p^\pi$, trained with the same architecture as a δ-model but with$\gamma = 0$and a single model atom. Return distributions are estimated by sampling full trajectory rollouts from the learned model and computing discounted returns on the sampled rewards. - Mean (Dirac at MC expected return): a degenerate return distribution consisting of a Dirac mass centered at the true expected return estimated by Monte Carlo—representing the best possible prediction that uses only first-moment information.
- γ-Model Ensemble with variance imputation: not described in detail but listed in Figure 3b as constructing return distributions from the ensemble's value function predictions.
- γ-Model Ensemble: an ensemble of
-
Generation budget / compute accounting. The paper does not use a standardized "generation budget" in the same way as the search-and-revisions paper. Instead, the δ-model is trained once on trajectory data from a fixed policy, and evaluation is zero-shot: return distributions for novel reward functions are produced by sampling model atoms from the learned δ-model and plugging the reward function into Equation 4, without any additional environment interaction. The cost of this post-hoc evaluation is the number of model samples drawn to build the empirical return distribution—not directly reported but implicitly controlled by the number of model atoms
$m = 51$(Pendulum) and the number of state samples per atom$s = 32$. Training cost is$3 \times 10^6$gradient updates for Pendulum and$1 \times 10^6$for Windy Gridworld. -
Cross-validation / statistical protocol. No cross-validation is reported. The paper evaluates on 9 source states for Pendulum and a single source state (the origin) for the Windy Gridworld visualizations. The MC ground truth return distributions are estimated via environment rollouts using the true policy and environment dynamics, providing an empirical reference for the predicted distributions. The reported Wasserstein distances aggregate across the 9 source states and (for some results) across the four held-out reward functions.
Main Quantitative Results
Zero-Shot Return Distribution Prediction Quality (Figure 3b, Pendulum)
The central quantitative result compares the δ-model's zero-shot return distribution predictions against all baselines on the Pendulum environment across four held-out reward functions (Default, Above Horizon, Stay Left, Counterclockwise Penalty; described in Appendix D.3.2).
Figure 3b reports the Cramér distance between each method's predicted return distributions and the MC ground truth, averaged over 9 source states:
-
δ-model (DSM) achieves the lowest Cramér distance across all four reward functions, substantially outperforming all baselines. The exact numerical values are not reported in the text, but the relative ordering is clear from the bar chart in Figure 3b: DSM < γ-Model Ensemble < Rollout < Mean, with the margin between DSM and the next-best baseline (γ-Model Ensemble) being substantial for the Default and Above Horizon rewards.
-
Rollout (learned
$p^\pi$) performs worst or second-worst on all reward functions, consistent with the paper's claim that:
"accumulation of error throughout sampled trajectories prevents this model from achieving reasonable return distributions, which is consistent with the difficulties of accurately rolling out long trajectories from learned forward models (Jafferjee et al., 2020; Abbas et al., 2020; Lambert et al., 2022)."
- γ-Model Ensemble outperforms Rollout despite not modeling aleatoric uncertainty, because:
"its superior ability to model long-horizon behavior enables it to estimate return distributions more accurately, achieving similar quality to a Dirac mass centered at the ground truth mean return."
In other words, the γ-model ensemble's predictions, while capturing only epistemic uncertainty (which collapses toward the mean with ensemble size), are less corrupted by compounding model error than the learned transition model's rollout-based estimates—even though the latter is in principle capable of generating proper distributional estimates. This is a practical vindication of the paper's argument that successor-type models are more robust to estimation error for long-horizon prediction than one-step models.
- Mean (Dirac at MC return) serves as a surprisingly strong baseline, achieving similar or slightly better Cramér distance than the γ-model ensemble for some rewards. This indicates that for these tasks, simply predicting the correct mean return (which the DSM factorization also enables) already captures a substantial fraction of the distribution's location, and the remaining challenge is to accurately model higher-order moments and tail behavior—which is precisely where the DSM outperforms the Mean baseline.
Qualitative Visualization of Model Atoms (Figure 4a, Windy Gridworld)
Figure 4a provides a qualitative comparison between the δ-model's distributional SM and an ensemble of γ-models trained on data from a uniform random policy in the Windy Gridworld environment. The visualization shows contour plots of the occupancy distributions predicted by each method from a source state at the origin:
- DSM model atoms (top-left, colored densities): The
$m = 4$atoms exhibit clearly distinct occupancy patterns, with different atoms concentrating probability mass in different regions of the grid—reflecting the multimodal nature of the uniform random policy's behavior under stochastic wind. The paper states:
"Due to the nature of the wind in this domain, which always forces the agent to the corner of the quadrant where it is located, a uniform random policy exhibits a multimodal distribution of model atoms, as shown by the colored densities in the top-left."
- γ-Model Ensemble atoms (bottom-left): The ensemble members all predict similar occupancy distributions that closely align with the standard successor measure (the mean occupancy). The paper notes:
"when examining an ensemble of γ-models trained on the same data, we see that the models in the ensemble all predict similar state occupancies which align closely with the SM – crucially, only the distributional SM captures the diversity of 'futures' that the agent can experience."
- DSM mean (top-right) vs. SM (bottom-right): The average over DSM model atoms (top-right) closely matches the successor measure itself (bottom-right, shown as the standard Ψπ), confirming that the δ-model's first moment correctly recovers the standard SR. The key point is that the DSM captures both the mean structure (via the ensemble average) and the full distribution over possible futures (via the diversity of atoms), while the γ-model ensemble captures only the mean.
This visualization is the paper's primary evidence that the δ-model's coupled training (via the model MMD loss) produces meaningfully diverse atoms that correspond to genuinely different behavioral outcomes, rather than just noisy variations around the mean. An ensemble of independently-trained γ-models, in contrast, converges to similar solutions because each member optimizes the same objective independently—there is no incentive for diversity.
Risk-Sensitive Policy Selection (Figure 4b, Windy Gridworld)
The paper's most distinctive experimental result demonstrates a capability that no prior method can provide. Figure 4b shows that the DSM enables correct ranking of two policies ($\pi_1$, $\pi_2$) across two reward functions (Hopscotch and Lopsided Checkerboard) with respect to both mean return and 0.4-CVaR, without any additional data collection or training beyond what was used to learn the distributional SMs.
Experimental setup: Distributional SMs are trained for two different policies in the Windy Gridworld. The reward functions are held out—the DSM was trained without reward information, so all evaluations are zero-shot. For each policy-reward pair, the DSM predicts a return distribution via Equation 4, from which the mean and 0.4-CVaR are computed. These are compared to ground-truth values estimated by Monte Carlo environment rollouts.
Findings (Figure 4b, bar charts):
-
For the Hopscotch reward: Both the DSM and MC estimates agree that
$\pi_1$is superior under both mean and CVaR criteria (DSM and MC bars for$\pi_1$are higher than for$\pi_2$under both metrics). -
For the Lopsided Checkerboard reward: There is a reversal in the optimal policy depending on the risk criterion. Under mean return,
$\pi_1$is superior (higher solid blue bar for$\pi_1$than for$\pi_2$). Under 0.4-CVaR,$\pi_2$is superior (higher solid pink bar for$\pi_2$than for$\pi_1$). The DSM correctly identifies this reversal in both cases—the solid bars (DSM predictions) align with the dashed bars (MC ground truth) in direction and approximate magnitude.
The paper emphasizes the uniqueness of this result:
"We note that, to our knowledge, no other method can accomplish this feat. On the one hand, existing distributional RL algorithms could not evaluate the return distributions for held-out reward functions. On the other hand, any algorithm rooted in the SM for zero-shot evaluation can only rank policies by their mean returns, so they must fail to rank
$\pi_1$,$\pi_2$by at least one of the objectives."
Why this matters concretely: The Lopsided Checkerboard scenario models a realistic decision problem where a policy that performs well on average ($\pi_1$) has a worse tail risk (as measured by CVaR) than an alternative ($\pi_2$). A decision-maker using only the standard successor measure would select $\pi_1$ (higher mean) without ever knowing that $\pi_2$ provides better downside protection. The DSM reveals this tradeoff without requiring any new environment interaction—the risk criterion can be changed at evaluation time, and the correct ranking emerges from the same learned distributional SM.
Monte Carlo Approximation in a Three-State MDP (Figures 5 and 6, Appendix C.1)
As a pedagogical illustration, Appendix C.1 includes Monte Carlo approximations to the distributional SM in tabular environments, generated by sampling 1,000–10,000 trajectories per state and converting them into visitation distributions. Figure 5 shows kernel density estimates of the distributional SM over the 2-simplex for a three-state MDP with a specific transition matrix and $\gamma = 0.7$. Figure 6 shows that for an MDP with uniform transition probabilities and $\gamma = 0.5$, the distributional SM is supported on the Sierpiński triangle—a fractal distribution that is neither discrete nor absolutely continuous with respect to Lebesgue measure. This is noted as a higher-dimensional analog of the Cantor distribution described in Bellemare et al. (2023, Example 2.11) for scalar return distributions, and it illustrates that distributional SMs can exhibit complex, fractal structure even in simple finite-state MDPs.
These tabular examples are not part of the main experimental evaluation but serve to ground the abstract mathematical object in concrete, visualizable instances and to demonstrate the breadth of distributional forms that $\mathcal{M}^\pi$ can take.
Ablation Studies and Robustness Checks
All ablation studies are reported in Appendix E, with results visualized in Figure 9.
- n-step bootstrapping horizon (Figure 9, left): The paper sweeps
$n \in \{1, 2, ..., 10\}$on the Pendulum environment with$\gamma = 0.95$and all other hyperparameters fixed. The Wasserstein distance (averaged over 9 source states and 4 reward functions) decreases substantially as$n$increases from 1 to approximately 5, after which benefits level off. At$n = 5$, the bootstrap term accounts for$\gamma^5 \approx 77\%$of the target distribution's mass, consistent with the paper's claim that:
"training stability tends to improve substantially when bootstrap samples account for roughly 80% of the samples."
Using $n = 1$ (pure one-step bootstrap) yields noticeably worse Wasserstein distance, confirming that long-horizon prediction from bootstrapped targets is a practical challenge for geometric horizon models, and that multi-step returns are an effective mitigation.
- Number of model atoms (Figure 9, middle): The number of model atoms
$m$is swept across$\{5, 11, 21, 31, 41, 51\}$on Pendulum. Wasserstein distance monotonically decreases as$m$increases, with the improvement most pronounced when going from 5 to 21 atoms and continuing more gradually through 51 atoms. The paper notes:
"As expected we obtain a better approximation to
$\mathcal{M}^\pi$when increasing the number of model atoms. Further scaling the number of model atoms should continue to improve performance at the cost of compute. This is a desirable property for risk-sensitive applications where better approximations are required."
This monotonic improvement supports the particle approximation framework: more atoms provide a finer discretization of the distribution over occupancy measures, and there is no observed saturation point at $m = 51$.
-
Adaptive (adversarial) vs. non-adaptive state kernel (Figure 9, right): The δ-model is trained with and without the adversarial kernel (the iResNet-based learned embedding). With the non-adaptive kernel (a fixed mixture of rational quadratic kernels applied directly to raw states), the Wasserstein distance is roughly doubled compared to the adaptive version. The paper quantifies this with a bar chart in Figure 9 (right) but does not report exact numerical values. This is the most impactful single ablation: removing the adaptive kernel degrades performance more severely than any other single architectural choice, confirming that non-stationarity of the learned distributions makes fixed similarity measures inadequate, and that the adversarial kernel is not a minor optimization but a critical component for learning accurate distributional SMs.
-
γ-Model ensemble vs. DSM (Figure 4a, qualitative): While not a formal ablation, the comparison between the δ-model and an independently-trained γ-model ensemble serves as an ablation of the coupled training via model MMD. The γ-model ensemble uses the same architecture and the same number of atoms (
$m = 4$for Windy Gridworld) but trains each atom independently with a state-level MMD loss, without the model-level MMD that couples atoms through the distributional Bellman residual. The result is that γ-model atoms collapse to similar occupancy distributions (all approximating the mean successor measure), while δ-model atoms remain diverse and capture distinct behavioral modes. This demonstrates that the model MMD loss is the mechanism that induces and preserves diversity among atoms—without it, the equally-weighted particle representation would degenerate into$m$copies of approximately the same distribution, providing no distributional information beyond the mean.
Critical Assessment
Claim: The δ-model enables zero-shot distributional policy evaluation that substantially outperforms baselines.
What was tested: The paper compares the δ-model's return distribution predictions against a learned transition model (Rollout), an ensemble of γ-models, and a Dirac-at-the-mean baseline on Pendulum across four reward functions. The δ-model achieves the lowest Cramér distance in all cases (Figure 3b).
What this demonstrates: The experiment convincingly shows that the δ-model produces better return distribution estimates than a learned one-step model rolled out over long horizons, and better than an ensemble that captures only epistemic uncertainty. The poor performance of the rollout baseline is consistent with the paper's stated motivation that compounding model error makes direct transition-model-based distributional estimation unreliable.
Limitations:
-
No comparison to standard distributional RL methods trained per-reward: The baselines are all zero-shot methods. A natural comparison would be to a standard distributional RL algorithm (e.g., QR-DQN, Dabney et al., 2018b; or MMD-DQN, Nguyen-Tang et al., 2021) trained from scratch on each reward function with the same amount of environment interaction data used to train the δ-model. This would quantify the cost of zero-shot generality: how much accuracy is sacrificed relative to a reward-specific method that cannot transfer? The paper does not address this tradeoff.
-
Only deterministic policies are evaluated: The experiments use fixed policies with data collected from them. All the distributional structure captured by the DSM in these experiments is therefore due to environment stochasticity (wind in Windy Gridworld, dynamics in Pendulum). The more interesting case—a stochastic policy producing multimodal occupancy distributions through deliberate action randomization—is not tested. This limits the demonstration of the DSM's ability to capture policy-induced aleatoric uncertainty.
-
Small scale of evaluation: The Pendulum evaluation uses 9 source states and 4 reward functions. The Windy Gridworld uses a single source state (the origin) for qualitative visualizations and two policies for the risk-sensitive selection experiment. No confidence intervals, standard errors, or significance tests are reported. The sample sizes are small enough that the observed ranking of methods could be sensitive to the specific source states and reward functions chosen.
-
No evaluation on high-dimensional state spaces: Both environments are low-dimensional (2D Windy Gridworld, ~3D Pendulum state). The paper provides no evidence that the δ-model scales to the high-dimensional state spaces (e.g., image-based observations) where generative modeling of state occupancies would be most challenging. The iResNet kernel is specifically designed to handle injectivity in arbitrary dimensions, but this is not empirically tested.
-
The γ-model ensemble baseline may be weaker than necessary: The paper notes that the γ-model ensemble is "almost equivalent to a δ-model, with the difference being that the individual γ-models of the ensemble are trained independently rather than coupled through the model MMD loss." However, this baseline does not have access to the n-step bootstrapping or the adversarial kernel innovations that the δ-model benefits from, making it difficult to isolate whether the performance gap is due to the distributional SM concept itself or to these training techniques. A fairer baseline would be an ensemble of γ-models with the same n-step and adversarial kernel enhancements, differing only in the coupling loss.
Claim: The DSM enables zero-shot risk-sensitive policy selection that no prior method can provide.
What was tested: Figure 4b shows that for two policies and two reward functions in Windy Gridworld, the DSM correctly identifies which policy is superior under both mean and CVaR criteria, including a scenario where the ranking reverses between criteria.
What this demonstrates: This is an existence proof that the capability is achievable—the DSM can correctly rank policies by risk-sensitive criteria on held-out rewards. The reversal case (Lopsided Checkerboard, where $\pi_1$ wins on mean but $\pi_2$ wins on CVaR) is a compelling demonstration that the distributional information matters for decision-making.
Limitations:
-
Only two policies and two reward functions: This is a minimal demonstration. The paper does not show how ranking accuracy scales with the number of policies or reward functions, whether the DSM maintains correct rankings across a broader distribution of tasks, or what the failure modes are when rankings are incorrect.
-
No analysis of when the DSM fails at risk-sensitive ranking: The paper shows a success case but provides no systematic study of conditions under which the DSM might incorrectly rank policies—for example, when the return distributions are very close, when the policies are very similar, or when the reward function has structure that interacts poorly with the finite particle approximation.
-
CVaR at a single level (0.4) only: The paper mentions the importance of arbitrary risk criteria but only demonstrates one specific CVaR level. Whether the DSM's distributional estimates are accurate enough to support evaluation across a full spectrum of risk levels (e.g., CVaR at 0.1, 0.2, ..., 0.9) or for other functionals (variance, entropy, probability of exceeding a threshold) is not tested.
Claim: The δ-model captures aleatoric uncertainty that ensemble methods miss.
What was tested: Figure 4a shows that the DSM atoms exhibit diverse occupancy patterns corresponding to distinct behavioral modes, while the γ-model ensemble atoms all approximate the mean successor measure.
What this demonstrates: The qualitative differences in Figure 4a are visually compelling and support the conceptual distinction the paper makes between epistemic and aleatoric uncertainty. The DSM atoms are meaningfully different (corresponding to different corridors the agent might be blown into), while the γ-model ensemble atoms are nearly identical.
Limitations:
-
This is a qualitative result in a single environment with a single policy: There is no quantitative metric for atom diversity, no demonstration that the diversity increases or decreases appropriately with environment stochasticity (e.g., does the DSM produce more diverse atoms when wind variance is higher?), and no test in the Pendulum environment.
-
The coupling mechanism that produces diversity is not ablated: The paper attributes atom diversity to the model MMD loss, but the precise mechanism is not isolated. It is possible that the diversity emerges not from the two-level MMD structure per se, but from the adversarial training dynamic (the discriminator incentivizing the generators to cover different modes) or from random initialization combined with the n-step bootstrapping. An ablation where the model MMD loss is replaced with a sum of independent state MMD losses (but keeping the adversarial kernel and n-step bootstrapping) would directly test this claim, but is not reported.
Claim: The distributional Bellman operator is a contraction and dynamic programming converges.
What was tested: Proposition 3.6 and Corollary 3.7 are proven mathematically in Appendix B.1. No empirical validation of the contraction property (e.g., measuring the distance between successive iterates and verifying geometric decay at rate $\gamma$) is reported.
Limitations:
-
No empirical convergence diagnostics: The paper provides theoretical guarantees but does not show that the δ-model's training dynamics exhibit the predicted contraction behavior. Metrics like the model MMD loss over training time, or the distance between successive δ-model iterates, could empirically validate that the learned approximation is converging to a fixed point as the theory predicts. Without such diagnostics, a reader cannot assess whether the theoretical contraction survives the approximations introduced by finite particles, finite state samples, function approximation, and adversarial training.
-
The compactness condition for Corollary 3.7: The convergence proof requires
$(\mathcal{X}, d)$to be compact, which is satisfied for Windy Gridworld (bounded region) but may not hold for Pendulum depending on the state representation (angles are compact, angular velocities may not be explicitly bounded). The paper does not discuss whether this condition is satisfied in practice or whether violations could explain any training instability.
Missing Experiments That Would Strengthen the Paper
-
Scaling with state dimension: The environments used are low-dimensional. Testing on a higher-dimensional continuous control task (e.g., MuJoCo environments with state dimensionality ~10–20) would demonstrate whether the δ-model and its two-level MMD loss scale beyond toy problems.
-
Scaling with discount factor: All experiments use
$\gamma = 0.95$. The long-horizon prediction challenge becomes more acute as$\gamma \to 1$(effective horizon grows). Testing at$\gamma = 0.99$would stress-test the n-step bootstrapping approach and reveal whether the δ-model's advantage over rollout-based methods grows or shrinks with horizon. -
Comparison to a reward-specific distributional RL oracle: Training a standard distributional RL algorithm (with the same network capacity and data budget) separately on each reward function would quantify the "cost of generality"—how much accuracy is sacrificed to enable zero-shot transfer.
-
Empirical contraction verification: Tracking the model MMD loss and pairwise atom distances over training to verify that the learning dynamics reflect the theoretical contraction property.
-
Diversity metrics for model atoms: Quantitative measures of atom diversity (e.g., average pairwise MMD between atoms, coverage of the state space, entropy of the categorical distribution over atoms when conditioned on reaching specific regions) would convert the qualitative claim of "capturing diverse futures" into a testable, quantitative statement.
-
Broader risk-sensitive evaluation: Systematic evaluation of the DSM's return distribution quality across multiple risk levels (CVaR at multiple
$\alpha$, variance, interquartile range) and for a larger set of reward functions and policies would strengthen the claim that the DSM enables "arbitrary risk criteria."
6. Limitations and Trade-offs
The Deterministic Reward Assumption Prevents Application to Most Real-World MDPs
The assumption or constraint. Proposition 3.2—the factorization identity that enables zero-shot distributional policy evaluation—requires the reward function $r: \mathcal{X} \to \mathbb{R}$ to be deterministic. The paper states this explicitly:
"Perhaps surprisingly, our assumption of a deterministic reward function made in Proposition 3.2 is necessary for a linear factorization between reward functions and return distributions."
Appendix C.2 provides a concrete counterexample showing why stochastic rewards break the factorization: two different sequences of states can produce identical occupancy distributions but different return distributions when rewards at individual states are random. The issue is that the mapping from state sequences to occupancy distributions is not injective—multiple trajectories can yield the same empirical visitation frequencies—and when rewards are stochastic, these trajectories produce different return distributions despite sharing an occupancy measure. The paper notes that the framework can be extended to "cycle-less environments" where each state is visited at most once per trajectory, which includes finite-horizon MDPs, but this is a severe restriction.
The consequence. Most MDPs of practical interest are both stochastic in rewards and cyclic. Any environment where the agent can revisit states—which includes almost all continuous control, navigation, and manipulation tasks with $\gamma < 1$—falls outside the factorization guarantee if rewards have any noise. Even with deterministic environment dynamics, if the reward is a noisy function of state (e.g., sensor noise, partial observability, reward shaping with stochastic components), the DSM factorization loses its distributional correctness. A practitioner evaluating risk-sensitive criteria on a new reward function would obtain return distributions that are mathematically incorrect in the stochastic reward case, with no bound on the magnitude of the error. The paper provides no diagnostic for detecting when this degeneracy occurs, and no partial correction.
What evidence exists in the paper. Appendix C.2 gives a worked example with a four-state MDP showing that two state sequences can map to identical occupancy distributions but produce return distributions $\mathcal{N}(0, \gamma^2)$ versus $\mathcal{N}(0, \gamma^4 + \gamma^6)$ when state $x_1$ has stochastic $\mathcal{N}(0,1)$ rewards. This is a constructive proof that the factorization fails. The experimental domains (Windy Gridworld, Pendulum) both use deterministic reward functions, so the limitation does not appear in the empirical results. The paper does not measure how far the DSM predictions deviate from true return distributions as reward stochasticity increases.
Mitigation status. The paper acknowledges the limitation in a single remark (Remark 3.3) and Appendix C.2, but proposes no solution beyond noting that cycle-less environments are covered. This is a fundamental structural constraint of the approach, not an engineering limitation. The factorization $G^\pi_r(x) \stackrel{\mathcal{L}}{=} (1-\gamma)^{-1} \mathbb{E}_{X' \sim M^\pi(\cdot \mid x)}[r(X')]$ fundamentally assumes that the randomness in returns comes only from the randomness in state occupancy. There is no clear path to extending this to stochastic rewards without introducing a separate noise model for per-step reward variation, which would break the clean separation between transition structure and reward that is the paper's central motivation. Future work on this is not suggested.
The Difficulty Estimation Cost Is Not Accounted for in the Headline Zero-Shot Transfer Claim
The assumption or constraint. The δ-model must be trained on trajectory data from a fixed policy interacting with the environment. The paper frames zero-shot transfer as the ability to evaluate new reward functions without further data collection or training, and the experiments validate this by applying Equation 4 to held-out rewards. However, the δ-model itself requires a non-trivial training investment: $3 \times 10^6$ gradient updates for Pendulum ($m = 51$ atoms), $1 \times 10^6$ for Windy Gridworld ($m = 4$ atoms), using trajectory data collected under the policy of interest. This training cost, which includes generating environment interaction data regardless of whether rewards are observed, is analogous to the "difficulty estimation" cost in the reference paper's compute-optimal framework—it is a prerequisite expense that is not amortized in the zero-shot evaluation cost.
Moreover, the paper does not report the quantity of environment interaction data used for training. In standard distributional RL, the agent learns from reward-labeled transitions; the δ-model learns from reward-agnostic transitions, but it must still observe sufficient state-transition data to model the occupancy distribution. The data requirement may be comparable to or greater than training a reward-specific distributional RL agent, especially given that the δ-model must learn a more complex object (a distribution over occupancy distributions rather than a distribution over scalar returns).
The consequence. A practitioner comparing the δ-model to alternative approaches must weigh the upfront training cost against the downstream benefit of zero-shot reward evaluation. If the number of reward functions to be evaluated is small—say, 2 or 3—it may be cheaper to train a standard distributional RL agent separately on each reward function than to train a single δ-model that can evaluate all of them zero-shot. The paper provides no break-even analysis: how many reward functions must be evaluated before the δ-model's one-time training cost is amortized relative to per-reward distributional RL training? Without this, the practical value proposition of zero-shot transfer is unclear.
The δ-model also introduces a computational overhead at inference time that standard successor measures do not: to obtain a return distribution for a new reward function, one must sample multiple model atoms ($m = 51$ for Pendulum), draw state samples from each generative model, compute the expected reward under each sampled occupancy distribution, and aggregate the results. This is substantially more expensive than the standard successor measure's single expectation $\mathbb{E}_{X' \sim \Psi^\pi(\cdot \mid x)}[r(X')]$, albeit cheaper than running full environment rollouts.
What evidence exists in the paper. The paper reports training durations (3M updates, 1M updates) and atom counts but does not compare total environment samples or wall-clock time to alternative methods. The Pendulum experiments use $m = 51$ model atoms with $s = 32$ state samples per atom per update, meaning each gradient step evaluates $51 \times 32 = 1632$ state samples through the generator and discriminator networks—a non-trivial per-step cost that scales with $m \times s$. No experiment ablates the data efficiency of δ-model training relative to standard distributional RL or γ-model training.
Mitigation status. The paper does not address this tradeoff. The zero-shot capability is presented as an unqualified benefit. The training cost is reported as experimental hyperparameters but not framed as a limitation or compared to alternatives. Future work on reducing the data requirements or providing amortization analysis is not suggested.
The δ-Model Has Only Been Demonstrated on Low-Dimensional Continuous Environments with Simple Policies
The assumption or constraint. All experiments use two environments: Windy Gridworld (2D continuous state) and Pendulum (~3D continuous state: angle, angular velocity, and action-dependent dynamics). Both environments have smooth, low-dimensional state spaces well-suited to the MLP-based generator and iResNet kernel architectures. The policies evaluated are fixed and pre-specified (a uniform random policy for Windy Gridworld visualizations; unspecified policies for Pendulum and the risk-sensitive selection experiment). The paper provides no results on high-dimensional state spaces (e.g., images, joint angles of complex robots), stochastic policies with entropy regularization, or environments with discontinuous dynamics.
The iResNet kernel is theoretically capable of handling arbitrary input dimensions while maintaining injectivity, but the paper does not empirically validate that the two-level MMD loss remains effective when state samples are high-dimensional (where kernel-based two-sample tests are known to suffer from the curse of dimensionality; Gretton et al., 2012). The generative modeling of states via MLPs that take noise vectors as input also becomes more challenging as state dimension grows—mode collapse, poor sample diversity, and training instability are well-documented failure modes of GAN-like architectures in high dimensions.
The consequence. A practitioner considering the δ-model for a realistic application—robotic control with joint-space observations, autonomous driving with LIDAR inputs, game playing from pixels—has no evidence that the method scales. The central claim that the DSM captures "diverse futures" (aleatoric uncertainty over occupancy measures) has only been demonstrated in environments where the space of possible futures is small and easily separable (four wind-driven quadrants in Windy Gridworld). In a high-dimensional setting where occupancy distributions are supported on complex manifolds, the equally-weighted particle representation with $m = 51$ atoms may be grossly insufficient to cover the relevant modes, and the generative models may fail to produce high-fidelity state samples.
The reliance on fixed, pre-trained policies also limits the demonstrated utility. The paper does not show that the DSM can be trained online alongside policy improvement, or that the distributional information can guide policy optimization (as standard distributional RL does; Dabney et al., 2018a). The framework is presented purely for evaluation of existing policies, not for learning improved policies. Extending to policy optimization would require handling the non-stationarity of the occupancy distribution as the policy changes, which the current training procedure (fixed data collection from a fixed policy) does not address.
What evidence exists in the paper. The experiments are limited to what the paper describes. Figure 4a shows diverse atoms in the 2D Windy Gridworld. Figure 3b shows return distribution quality on the ~3D Pendulum. No results exist for higher-dimensional or more complex environments. The ablation on the number of model atoms (Figure 9, middle) shows monotonic improvement up to $m = 51$, but this is evaluated only on Pendulum and provides no evidence about how $m$ must scale with state dimension.
Mitigation status. The paper does not claim scalability to high-dimensional environments or to policy optimization. The scope is explicitly foundational: introduce the mathematical object, propose a tractable approximation, and demonstrate the core capability (zero-shot distributional evaluation) in proof-of-concept domains. The authors do not suggest specific approaches for scaling to higher dimensions, though the use of iResNet for the kernel and MLP generators could in principle be replaced with convolutional or transformer architectures suitable for image-based states. Future work on scaling is implied by the paper's positioning as a first step, but no concrete directions are proposed.
No Empirical Validation That Training Converges to the Theoretical Fixed Point
The assumption or constraint. Propositions 3.6 and Corollary 3.7 prove that the distributional Bellman operator $\mathcal{T}^\pi$ is a $\gamma$-contraction in the supremum-Wasserstein metric and that iterative dynamic programming converges to the unique fixed point $\mathcal{M}^\pi$. This theoretical guarantee is the foundation for the algorithm: if the operator did not contract, temporal-difference learning with function approximation would have no hope of converging to a meaningful solution.
However, the δ-model training deviates from the theoretical dynamic programming in several ways: (1) finite particle approximation ($m$ atoms instead of a continuous distribution), (2) finite state samples ($s = 32$ per generative model), (3) function approximation via neural networks (limited capacity, optimization error), (4) adversarial training (min-max rather than pure minimization), (5) $n$-step bootstrapping (which changes the effective operator), and (6) target network updates with Polyak averaging (introducing a moving target rather than fixed-point iteration). Whether the contraction property survives these approximations is not empirically verified.
The consequence. Without empirical convergence diagnostics, a practitioner cannot distinguish between a δ-model that has successfully approximated $\mathcal{M}^\pi$ and one that has converged to a poor local optimum of the adversarial training objective, or one that oscillates without converging. The theoretical guarantee provides motivation but no practical assurance. This is particularly concerning given the complexity of the training procedure—adversarial training with two-level MMD on a doubly-infinite-dimensional space—where issues like mode collapse, discriminator saturation, and training instability are well-documented challenges in GAN literature (Goodfellow et al., 2014; Binkowski et al., 2018).
Specific failure modes that could arise: the discriminator (kernel) could become too strong, making the model MMD insensitive to generator improvements and stalling learning; the generators could mode-collapse to a single occupancy distribution, losing all distributional information; the $n$-step bootstrapping could introduce bias that prevents convergence to $\mathcal{M}^\pi$ even if the training loss goes to zero; the Polyak-averaged target networks could lag too far behind the online networks, creating a moving target that prevents any fixed point from being reached.
What evidence exists in the paper. The paper reports no convergence diagnostics. There are no plots of the model MMD loss over training time, no Wasserstein distance to MC ground truth evaluated at intermediate checkpoints, no measurement of whether the distance between successive δ-model iterates decays at the theoretical rate $\gamma$. The training hyperparameters (3M gradient steps for Pendulum, 1M for Windy Gridworld) are stated as fixed budgets without evidence that training has converged by these points or that further training would not improve or degrade performance.
Mitigation status. Not addressed. The paper treats the theoretical contraction as sufficient justification for the learning algorithm and provides no empirical validation that the implemented system inherits the theoretical convergence properties. A simple diagnostic—tracking the model MMD loss during training and verifying that it decreases and stabilizes—would partially address this, but is not included. Future work on understanding the interaction between function approximation, adversarial training, and distributional dynamic programming in this nested MMD setting is not suggested.
The Risk-Sensitive Policy Selection Demonstration Is Minimal and Lacks Systematic Characterization of Failure Modes
The assumption or constraint. Figure 4b demonstrates that the DSM correctly ranks two policies on two reward functions under two risk criteria (mean and 0.4-CVaR) in the Windy Gridworld. This is the central experimental claim for the paper's unique capability: zero-shot risk-sensitive policy selection. However, the demonstration is extremely minimal: 2 policies $\times$ 2 rewards $\times$ 2 criteria = 8 total ranking decisions, with ground truth provided by Monte Carlo environment rollouts. The paper does not characterize how ranking accuracy varies with the number of policies, the similarity of the policies being compared, the structure of the reward function, the choice of risk level $\alpha$, or the environment stochasticity.
The consequence. A practitioner deploying the DSM for risk-sensitive decision-making needs to know not just that the method can work (as shown by the existence proof in Figure 4b) but when it works reliably. Several plausible failure modes are unexamined:
-
When return distributions overlap heavily: If two policies produce return distributions that are close (e.g., overlapping by 90% in probability mass), small estimation errors in the DSM could easily flip the ranking, particularly for tail-sensitive criteria like CVaR that depend on the extreme quantiles of the distribution. The finite particle representation with
$m = 4$atoms (Windy Gridworld) may not provide sufficient resolution to distinguish closely-matched policies. -
When the risk criterion depends on distributional features not well-captured by
$m$atoms: CVaR at level$\alpha$depends on the$\alpha$-tail of the return distribution. For small$\alpha$(e.g., CVaR 0.01), the estimate depends on the most extreme 1% of outcomes, which could be very poorly represented by$m = 4$or even$m = 51$atoms. -
When rewards have high dynamic range: If the reward function contains extreme values (very large positive or negative rewards in certain states), the plug-in estimator
$(1-\gamma)^{-1}\mathbb{E}_{X' \sim \theta_i(x)}[r(X')]$can produce high-variance estimates when the occupancy distribution$\theta_i(x)$has low but non-zero mass on the extreme-reward states. This variance could dominate the return distribution estimate. -
When policy differences are subtle: The paper does not characterize the two policies
$\pi_1$,$\pi_2$used in Figure 4b. If they are very different (e.g., one goes left, one goes right), the correct ranking is an easier test than if they differ only slightly in their action probabilities.
What evidence exists in the paper. Only Figure 4b, which shows 8 ranking decisions with no error bars, no statistical tests, and no characterization of the policies or the uncertainty in the MC ground truth rankings. The paper provides no sweep over risk levels $\alpha$, no evaluation of ranking accuracy across many reward functions, and no analysis of what happens when the DSM's ranking disagrees with MC (such cases may not appear in the minimal test set).
Mitigation status. The paper acknowledges none of these limitations. The risk-sensitive selection result is presented as a successful demonstration of capability, and the paper does not discuss the reliability, failure modes, or scope of validity of this capability. The ablation on the number of model atoms (Figure 9, middle) suggests that accuracy improves with $m$, but this ablation is on Pendulum return distribution quality (not on policy ranking accuracy) and does not address the specific challenges of risk-sensitive ranking. Future work on systematic evaluation of risk-sensitive policy selection across broader policy and reward distributions is not explicitly suggested but is a natural next step.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper introduces a new category of object into reinforcement learning—the distributional successor measure—and in doing so, fundamentally expands what "successor-type" representations can express. Prior to this work, the successor representation and its continuous analogs occupied a clear but limited niche: they factorized expected returns, enabling zero-shot policy evaluation across reward functions, but were silent on variability, risk, and tail behavior. Distributional RL, in turn, captured these distributional properties but was shackled to specific reward functions, requiring retraining for every new evaluation criterion. These two capabilities—zero-shot reward transfer and distributional uncertainty quantification—existed in separate methodological lineages with no bridge between them.
The distributional SM provides that bridge. It is not a compromise that weakly achieves both goals; it is a precise mathematical factorization (Proposition 3.2) showing that return distributions decompose into a reward-agnostic distribution over occupancy measures and the reward function itself. This factorization is the distributional analog of the identity $V^\pi_r = (1-\gamma)^{-1}\mathbb{E}_{X'\sim\Psi^\pi}[r(X')]$ that Dayan (1993) established for expected returns, and it inherits the same conceptual elegance: learn the environment-policy interaction once, then evaluate arbitrary rewards post-hoc. But unlike the standard SR, the distributional version preserves the full distributional information needed for risk-sensitive decision-making.
The magnitude of the shift is best characterized as a reframing with new capabilities rather than a paradigm shift. It does not overturn existing distributional RL or successor representation methods—both remain valid and useful in their original domains. Rather, it reveals that these two lines of work were studying different projections of the same underlying object. The standard successor measure $\Psi^\pi$ is the expectation of $M^\pi$; standard return distributions $\eta^\pi_r$ are the pushforward of $\mathcal{M}^\pi$ through the reward evaluation functional. The distributional SM sits at the intersection, and the paper's primary contribution is to name it, characterize it theoretically, and show it can be tractably approximated.
The paper resolves a latent contradiction in the literature that was less visible than the one in the reference paper but equally real: why have distributional successor features (Gimelfarb et al., 2021; Carvalho et al., 2023) not been more successful at enabling risk-sensitive transfer? The answer, revealed by this paper's framework, is that these methods model marginal distributions over scalar features rather than the joint distribution over state occupancies. The distributional SM makes clear that the right distributional object is not a distribution over the output of the SR (a scalar), but a distribution over the input to the SR's linear operator (an occupancy measure). Modeling marginals discards the correlation structure across states that is essential for correctly computing return distributions for arbitrary reward functions. This diagnosis is constructive: it tells future researchers what they need to model (joint distributions over occupancies), not just that previous attempts were insufficient.
This work also makes the research direction of learning distributional forward models more attractive while making scalar-feature distributional transfer less attractive. The paper's empirical demonstration that a learned transition model $\hat{p}^\pi$ fails badly at zero-shot return distribution estimation due to compounding error (Figure 3b, "Rollout" baseline) is not merely a negative result—it is a constructive argument for why temporally-abstracted distributional models are necessary. The distributional SM does not avoid the transition model entirely (Proposition F.1 shows $\mathcal{M}^\pi$ is determined by $p^\pi$), but it learns to answer long-horizon queries directly rather than through repeated one-step prediction. This temporal abstraction paradigm—learning to predict discounted sums rather than one-step transitions—is validated as being more robust to model error in a distributional setting, extending the known advantages of successor measures over one-step models (Janner et al., 2020; Thakoor et al., 2022) to the full distributional case.
The paper also elevates the status of MMD as a distributional RL loss, following Nguyen-Tang et al. (2021) and Zhang et al. (2021b) but extending it to the nested $\mathcal{P}(\mathcal{P}(\mathcal{X}))$ setting. The demonstration that a two-level MMD construction—with a characteristic model kernel built from a state-level MMD—provides a practical and theoretically-grounded loss on this complex space is a methodological contribution that may generalize. Any problem involving distributions over structured objects where a kernel can be defined on the base space can potentially adopt this compositional MMD approach, making it a design pattern rather than a one-off solution.
Finally, the paper's risk-sensitive policy selection result (Figure 4b) serves as an existence proof that zero-shot distributional evaluation can drive meaningful decisions. It is a single, small-scale demonstration, but it concretely shows something that no prior method could do: correctly identify that a policy which is better on average can be worse on a tail-risk measure, on a reward function never seen during training. This capability is not just a benchmark improvement—it is a qualitatively new tool for decision-making under uncertainty that decouples the expensive part (learning environment dynamics) from the cheap part (evaluating reward functions and risk criteria). The practical implications of this decoupling are substantial, even if the current implementation is limited to low-dimensional environments with deterministic rewards.
Follow-Up Research This Work Enables
Scaling the δ-model to high-dimensional state spaces with image-based observations. The current experiments use 2D and ~3D continuous state spaces where MLP generators and iResNet kernels suffice. The central question is whether the two-level MMD loss and equally-weighted particle representation remain effective when states are images (e.g., 84×84×3 pixels) and occupancy distributions are supported on complex manifolds. A concrete experiment: train a δ-model on Atari environments (using a fixed policy for evaluation, as in the paper) with convolutional generators (DCGAN-style or StyleGAN-based) and a convolutional iResNet for the adversarial kernel. Measure whether the model atoms exhibit meaningful diversity (do different atoms correspond to meaningfully different strategies or outcomes, such as reaching different rooms in Montezuma's Revenge?) or whether they collapse to minor pixel-level variations that are perceptually identical. The key metric would be a feature-space diversity measure—e.g., clustering state samples from different atoms in the embedding space of a pretrained dynamics model—to distinguish genuine behavioral diversity from low-level noise. This experiment would directly test whether the distributional SM framework scales beyond the toy domains where it has been demonstrated, and negative results (atoms collapsing despite the model MMD loss) would identify fundamental limitations of the particle representation for high-dimensional occupancy distributions.
Stochastic reward extension via trajectory partitioning. The deterministic reward assumption (Proposition 3.2 and Appendix C.2) is the most significant theoretical limitation of the current framework. Appendix C.2 shows that the factorization fails because multiple trajectories can map to the same occupancy distribution but produce different return distributions under stochastic rewards. A natural direction is to learn a partition of the trajectory space such that within each partition element, the mapping from trajectories to occupancy distributions is approximately injective. Concretely, augment the δ-model to output not just occupancy distributions but also a latent variable that disambiguates trajectories with identical occupancies but different reward sequences. The training objective would encourage this latent variable to be predictive of per-step rewards (when rewards are available during training) while the occupancy atoms remain reward-agnostic. At inference time on a new reward function, one would need to impute the latent variable—perhaps via a learned prior or by marginalizing over it. A strong follow-up would test this on a gridworld with stochastic rewards, comparing the augmented DSM's return distribution predictions against both the deterministic-reward DSM (which would be biased) and a reward-specific distributional RL oracle (which would be unbiased but not zero-shot). The key result would be whether the augmented model achieves near-oracle return distribution accuracy while retaining zero-shot capability. If the latent variable approach fails—if the required number of latent categories scales poorly with environment complexity—this would suggest that the distributional SM factorization has fundamental limits in stochastic-reward settings, which is itself an important negative result to establish.
Online distributional SM learning with policy improvement. The paper treats the δ-model purely as a policy evaluation tool: a fixed policy is used to collect data, the DSM is learned, and then policies are evaluated post-hoc. A natural extension is to integrate the DSM into an actor-critic or policy iteration loop where the policy improves over time. The central challenge is that as the policy changes, the occupancy distribution $\mathcal{M}^\pi$ changes, creating a non-stationary target for the δ-model. A concrete experiment: implement a distributional policy iteration algorithm where (1) the current policy's DSM is estimated from recent trajectory data using the δ-model training procedure, (2) for a set of candidate reward functions (representing different objectives or constraints), the DSM is used to compute risk-sensitive return estimates, and (3) the policy is updated to improve a composite objective (e.g., maximize mean return on task A subject to a CVaR constraint on task B). The key measurement would be whether the DSM can track the changing policy quickly enough to guide improvement, or whether the temporal-difference learning of $\mathcal{M}^\pi$ lags behind policy changes, leading to stale distributional estimates that mislead the policy update. A positive result would demonstrate that the distributional SM framework extends from evaluation to optimization, making it a full replacement for standard distributional RL in multi-task settings. A negative result—the DSM learning is too slow to keep up with policy changes, and the resulting policy improvement is worse than training separate distributional RL agents per task—would clarify that the distributional SM is best suited for post-hoc analysis of fixed policies rather than online learning.
Diagnosing and mitigating the curse of dimensionality in two-level MMD. Kernel two-sample tests, including MMD, are known to suffer from the curse of dimensionality: as the data dimension grows, the test power degrades unless the kernel bandwidth is carefully adapted or the data lies on a lower-dimensional manifold (Gretton et al., 2012). The δ-model compounds this issue by using MMD at two nested levels: state-level MMD (on $\mathcal{X}$) and model-level MMD (on $\mathcal{P}(\mathcal{X})$, via the state MMD). A systematic study of how the model MMD's discriminative power scales with state dimension would be valuable. A concrete experiment: train δ-models on environments with controlled state dimensionality (e.g., $d \in \{2, 4, 8, 16, 32, 64\}$) while keeping the environment dynamics (e.g., linear-Gaussian transitions) constant. Measure (a) the convergence rate of the model MMD loss, (b) the diversity of learned atoms (average pairwise state MMD), and (c) the accuracy of zero-shot return distribution predictions as a function of $d$. If performance degrades sharply beyond some threshold dimension, investigate mitigations: (1) replacing the iResNet with a kernel that explicitly performs dimensionality reduction while preserving injectivity (e.g., a variational information bottleneck), (2) using sliced MMD or random Fourier features to reduce the effective dimension, or (3) learning a lower-dimensional latent dynamics model and defining the state kernel in latent space. This study would provide practical guidance on the dimensionality regimes where δ-models are viable and identify architectural innovations needed to push that boundary.
Verifier over-optimization in distributional successor measures. The reference paper on compute-optimal test-time scaling identified verifier over-optimization as a central bottleneck: aggressive optimization against a learned verifier (process reward model) causes performance degradation because the verifier's scores are imperfect. The distributional SM faces a structurally similar challenge: the δ-model is optimized adversarially against a learned kernel (the discriminator), and the quality of the learned distributional SM depends on the discriminator's ability to provide accurate MMD estimates. If the discriminator overfits or becomes poorly calibrated, the δ-model may learn to exploit these imperfections—producing occupancy distributions that score well under the model MMD but are poor approximations to $\mathcal{M}^\pi$. A concrete experiment: track both the adversarial training loss (generator vs. discriminator) and the ground-truth Wasserstein distance between the DSM predictions and MC return distributions during training. Look for a "turnaround point" where the adversarial loss continues to improve (generator "fooling" the discriminator) but the ground-truth distance begins to increase. If such a point exists, it would indicate verifier over-optimization in the distributional SM setting, analogous to the beam search degradation observed in the reference paper. Mitigations to test: early stopping based on a validation metric, adding noise to discriminator inputs, using an ensemble of discriminators, or applying a gradient penalty to constrain the discriminator's Lipschitz constant (as in WGAN-GP). This experiment would connect the distributional SM literature to the broader discussion of over-optimization in learned verifiers and provide practical training guidelines.
Combining distributional SMs with successor features for generalized policy improvement. The successor features (SF) framework (Barreto et al., 2017) extends the successor representation to the case where rewards are linear in a known feature map $\phi: \mathcal{X} \to \mathbb{R}^d$, enabling generalized policy improvement across tasks: given successor features for multiple policies, one can construct a new policy that outperforms all of them on a new task, without further learning. The distributional SM provides the distributional analog of the successor measure, but does not address the policy improvement side. A natural synthesis is to define distributional successor features: learn, for each policy, a distribution over discounted feature occupancies $\sum_t \gamma^t \phi(X_t)$ rather than just a distribution over state occupancies. This would combine the zero-shot distributional evaluation capability of the DSM (evaluating any reward function expressible as a linear combination of features) with the policy improvement capability of SFs. A concrete experiment: train distributional SFs for a set of policies in a continuous control task where rewards are linear combinations of features (e.g., position, velocity, energy). For a new reward weighting, use the distributional SFs to (a) compute the return distribution for each existing policy (zero-shot distributional evaluation), (b) compare policies under risk-sensitive criteria (as in Figure 4b), and (c) attempt generalized policy improvement by constructing a new policy from the per-state distributional information. The key question is whether the distributional information (particularly the joint distribution over features) enables better policy improvement than standard SFs, which only use expected feature occupancies. If distributional SFs yield better policies under risk-sensitive objectives (e.g., lower CVaR), this would establish that distributional information is not just useful for evaluation but for optimization as well.
Practical Applications and Downstream Use Cases
Post-deployment safety auditing of learned policies without additional environment interaction. Consider a robotics company that has trained a navigation policy for a warehouse robot using a standard RL objective (minimize delivery time). Before deployment, the safety team needs to evaluate this policy against multiple failure modes: collision risk with shelving units (reward: -1000 per collision), proximity to human workers (reward: -10 per second within 1 meter), and energy consumption (reward: -0.1 per unit of torque). With standard distributional RL, each of these evaluations would require training a separate return distribution predictor from scratch, which may require deploying the robot with untested safety constraints—an unacceptable risk. With a pre-trained distributional SM for the navigation policy (trained once on state-transition data from the robot's standard operation, without any of these safety reward functions), the safety team can compute the full return distribution for each failure mode in seconds by plugging the corresponding reward function into Equation 4. They can then assess not just the expected collision rate, but the distribution—including the probability of catastrophic outcomes (e.g., the 99th percentile of collision severity). The paper's demonstration in Figure 4b, while small-scale, shows that this zero-shot risk-sensitive evaluation can correctly rank policies by both mean and CVaR on held-out rewards. The key practical requirement is that the reward functions be deterministic functions of state, which covers many safety-relevant metrics (collision = boolean function of proximity, energy = deterministic function of torque). For a warehouse robot operating in a ~10-50 dimensional state space (joint angles, velocities, LIDAR readings), extending the δ-model to this dimensionality is the primary engineering challenge, and the proposed scaling experiments in the previous section would directly inform feasibility.
Multi-stakeholder decision-making for autonomous systems with evolving preferences. An autonomous vehicle company develops a driving policy optimized for a base reward function (progress toward destination, lane-keeping, speed maintenance). After deployment, regulators introduce new requirements: limits on hard braking events, pedestrian proximity buffers that vary by neighborhood, and emissions constraints that depend on local air quality indices. Each of these corresponds to a different reward function over the vehicle's state space, and the relative importance of these criteria may change over time (e.g., stricter emissions in summer months). Retraining the distributional RL system for each combination of stakeholder preferences is impractical—there are combinatorially many possible weightings. The distributional SM offers an alternative: learn $\mathcal{M}^\pi$ once for the base driving policy, and then at decision time, aggregate the return distributions for different reward components into a composite risk-adjusted score. The vehicle does not need to re-learn how it drives; it only needs to re-evaluate what it is optimizing for. The paper's zero-shot evaluation capability (Figure 3a-b, four reward functions evaluated without retraining on Pendulum) demonstrates the core mechanism. The practical challenge is that real-world driving states are high-dimensional (camera images, sensor fusion) and the policy may need to adapt online to new reward weightings, requiring fast DSM sampling (Equation 4) at inference time. The per-inference cost of drawing $m = 51$ model atoms and $s = 32$ state samples each is non-trivial but could be amortized by pre-computing occupancy samples for common origin states. This use case also motivates the stochastic reward extension discussed above, since sensor noise could make rewards like "pedestrian proximity" stochastic in practice.
Efficient simulation-based design optimization with learned surrogate distributional models. An engineer is designing a controller for a wind turbine across multiple operating conditions (wind speeds, turbulence levels). For each candidate controller design, they need to evaluate performance across multiple criteria: energy output, mechanical stress on the gearbox, and noise emissions. Running high-fidelity fluid dynamics simulations to evaluate each controller on each criterion is computationally prohibitive—each simulation takes hours. The distributional SM provides a surrogate model: train a δ-model on simulation data from a few representative wind conditions (capturing the distribution of the turbine's state trajectory under a given controller), and then evaluate any reward function (energy, stress, noise) zero-shot via Equation 4. Moreover, because the DSM captures the full distribution of state occupancies, the engineer can assess not just expected performance but also risk measures: the probability that gearbox stress exceeds the fatigue limit, or the 95th percentile of noise emissions during a gust. The paper's Figure 3b demonstrates that the DSM provides more accurate return distribution estimates than alternatives like learned one-step models with rollouts—specifically, the Rollout baseline fails due to compounding error, while the DSM maintains accuracy by learning temporally-abstracted occupancy distributions directly. The practical translation: for simulation-based design with expensive forward models, training a single DSM per controller design and using it as a multi-objective surrogate is more efficient than running separate simulations per objective, especially when the number of objectives is large or changes during the design process. The current limitation to low-dimensional state spaces means this is most directly applicable to engineering domains where the state can be adequately summarized by a modest number of physical variables (position, velocity, stress, temperature) rather than full-field simulations (which would require scaling the δ-model to high-dimensional outputs, an open research question).
When to Prefer This Method
The paper explicitly positions the distributional SM against two classes of alternatives: (1) standard distributional RL methods that require reward-specific training, and (2) successor measure methods (including ensembles of $\gamma$-models) that enable zero-shot evaluation but only for expected returns. It also implicitly positions against learned transition models that could in principle provide zero-shot distributional evaluation through rollout sampling but fail in practice due to compounding error. The following decision rules emerge from the paper's theoretical and empirical analysis:
Prefer the distributional SM (δ-model) when:
- You need to evaluate a fixed policy against multiple reward functions without additional data collection, and at least some of these reward functions require distributional information (risk measures, uncertainty quantification) beyond the expected return. The paper demonstrates this on 4 held-out reward functions (Figure 3a-b) and 2 reward functions with risk-sensitive criteria (Figure 4b).
- The reward functions are deterministic (functions of state only). The factorization in Proposition 3.2 requires this, and Appendix C.2 provides a counterexample where stochastic rewards break the distributional correctness. If your rewards have stochastic components, the DSM will produce biased return distribution estimates with unknown error magnitude.
- The environment is low-to-moderate dimensional (the paper demonstrates 2D-3D state spaces). Scaling to high-dimensional states is plausible but empirically unvalidated. The iResNet kernel architecture is designed to handle arbitrary dimensions, but the two-level MMD loss may suffer from the curse of dimensionality in practice.
- You can afford the upfront training cost (3M gradient steps for Pendulum with 51 atoms, 1M for Windy Gridworld with 4 atoms) and will amortize it over many reward function evaluations. The paper does not provide a break-even analysis, but the cost is fixed regardless of how many reward functions are subsequently evaluated.
Prefer standard distributional RL (e.g., QR-DQN, MMD-DQN) when:
- You only need to evaluate a single reward function, or a very small number (1-3) of reward functions. In this regime, the upfront cost of training a δ-model likely exceeds the cost of training separate distributional RL agents per reward, especially since the δ-model must model
$\mathcal{P}(\mathcal{P}(\mathcal{X}))$rather than the simpler$\mathcal{P}(\mathbb{R})$. - The reward function is stochastic. Standard distributional RL handles stochastic rewards naturally by modeling the full randomness of the return, including both transition noise and reward noise, without requiring the injectivity condition that the DSM factorization needs.
- The policy is being actively improved (online RL), and the distributional return estimates need to track a changing policy. The paper does not address online learning, and the δ-model training procedure (fixed data collection from a fixed policy) would need significant modification to handle non-stationary occupancy distributions.
Prefer an ensemble of γ-models when:
- You only need expected returns (not distributions) for zero-shot reward evaluation. The γ-model ensemble is simpler to train (independent training per atom, no adversarial kernel, no model-level MMD) and provides good first-moment estimates. The paper shows that the γ-model ensemble's return distribution estimates are roughly comparable to placing a Dirac at the correct mean (Figure 3b, "Mean" vs. "γ-Model Ensemble"), so the distributional information it provides is minimal.
- You are in a data-limited regime where the coupled training of δ-model atoms (which requires the model MMD to maintain diversity) may be sample-inefficient relative to independent training. The paper provides no data-efficiency comparison, but adversarial training is generally more data-hungry than supervised or moment-matching objectives.
Do not prefer learned transition models with rollout-based distributional estimation when:
- The effective horizon
$(1-\gamma)^{-1}$is long. The paper's experiments with$\gamma = 0.95$(effective horizon ~20 steps) show the Rollout baseline performs substantially worse than the DSM and even worse than a simple Dirac at the mean (Figure 3b). The compounding error problem is well-documented and the paper's results are consistent with prior findings (Jafferjee et al., 2020; Lambert et al., 2022). For discount factors closer to 0.99 (effective horizon ~100), this failure mode would be exacerbated.
