Distributional value gradients for stochastic environments
Authors: Baptiste Debes
Tinne Tuytelaars
Venue: ICLR 2026 Poster (2026)
Keywords: Distributional Reinforcement Learning
Value Gradients
Sobolev Training
Stochastic Environments
MuJoCo Benchmarks
Noisy Dynamics
Primary area: reinforcement learning
TL;DR: We introduce Distributional Sobolev Training, which models distributions over values and their gradients via a Sobolev Temporal Difference operator, proving contraction and improving RL under stochastic dynamics.

Abstract: Gradient-regularized value learning methods improve sample efficiency by leveraging learned models of transition dynamics and rewards to estimate return gradients. However, existing approaches, such as MAGE, struggle in stochastic or noisy environments,  limiting their applicability. In this work, we address these limitations by extending distributional reinforcement learning on continuous state-action spaces to model not only the distribution over scalar state-action value functions but also over their gradients. We refer to this approach as Distributional Sobolev Training. Inspired by Stochastic Value Gradients (SVG), our method utilizes a one-step world model of reward and transition distributions implemented via a conditional Variational Autoencoder (cVAE). The proposed framework is sample-based and employs Max-sliced Maximum Mean Discrepancy (MSMMD) to instantiate the distributional Bellman operator. We prove that the Sobolev-augmented Bellman operator is a contraction with a unique fixed point, and highlight a fundamental smoothness trade-off underlying contraction in gradient-aware RL. To validate our method, we first showcase its effectiveness on a simple stochastic reinforcement‐learning toy problem, then benchmark its performance on several MuJoCo environments.

Published as a conference paper at ICLR 2026
DISTRIBUTIONAL VALUE GRADIENTS FOR STOCHAS- TIC ENVIRONMENTS
Baptiste Debes
Tinne Tuytelaars
PSI
PSI
KU Leuven
KU Leuven
Leuven, Belgium
Leuven, Belgium
baptiste debes @ebesskuleuven.br
tinne. tuyrellaasskkuuleeennn be
ABSTRACT
Gradient-regularized value learning methods improve sample efficiency by lever- aging learned models of transition dynamics and rewards to estimate return gra- dients. However, existing approaches, such as MAGE, struggle in stochastic or noisy environments, limiting their applicability. In this work, we address these limitations by extending distributional reinforcement learning on continuous state- action spaces to model not only the distribution over scalar state-action value func- tions but also over their gradients. We refer to this approach Distributional Sobolev Training. Inspired by Stochastic Value Gradients (SVG), our method uti- lizes one-step world model of reward and transition distributions implemented via a conditional Variational Autoencoder (cVAE). The proposed framework is sample-based and employs Max-sliced Maximum Mean Discrepancy (MSMMD) to instantiate the distributional Bellman operator. We prove that the Sobolev- augmented Bellman operator is contraction with unique fixed point, and highlight a fundamental smoothness trade-off underlying contraction in gradient- aware RL. To validate our method, we first showcase its effectiveness on simple stochastic reinforcement-leanning toy problem, then benchmark its performance on several MuJoCo environments
INTRODUCTION
Reinforcement learning (RL) tackles sequential decision-making by training agents maximize cumulative rewards. Off-policy actor-critic algorithms pair an actor, generating the control policy, with critic, estimating expected returns (i.e., the Q-function). This mapping from state-action pairs to expected returns, known as credit assignment, is typically learned via temporal-differee (TD) methods (Sutton 1988) and is critical for policy optimization in continuous-action settings. This paper is motivated by two lines of work aimed at improving credit assignment:
The use of action-gradients: In value-based continuous control, the critic's value is used directly to select actions but to provide aition-graddents for policy optimization (Lillicrap et al., 2016 Fujimoto et al., 20188 HHarmoja et al., 2018 D'Oro & Jaskowski (020). Conventional TD learning implicitly learns these gradients via value prediction, but this relies on smoothness assumptions the true value function that can degrade performance. To address this limitation,
DrOOr & Jaskowski (2020); Garibbo et al.) (2044) incorporate gradient information (Czameeki et a 2017) into critic training by learning a transition-reward world model (i.c. diffcrentiable proxy for the environment) and backpropagating through it (Heess et al. 2015).
Distributional RL: Many environments exhibit irreducible uncertainty in transitions and rewards.
Distributional RL (Morimura et al. 2010 Bellemare al. 2017 (023) captures this by modeling the return distribution rather than just its expectation. Categorical (Barth-Maron et al. 2018) and quantile-based (Dabney al. 2018b al approaches have provided rich and stable learning signals, yielding performance gains in variety of tasks (Barth-Maree et al. 2018 Dabney et al. 201 8a
Hessel al. 2017).
'The JAX Bradhury al. 2018 implementation is available at ht.pa://github..cm
Rapt wet.iistelebbeeebeeeee onal lue -ue--erdiii enta

---

Published as a conference paper at ICLR 2026
We argue that randomness also affects action gradients of returns, which have detrimental effect, especially in nigh-dimensiinal action spaces. As our experiments show (Section 6). existing methods that use gradient information deterministically (Czarnecki et al. 2017 D'Oro & Jaskowski, 2020) may struggle once the gradient to model becomes noisy or stochastic, losing some of the sample-effiiiecyy benefits of gradient modeling.
Paper contributions We extend distributional modeling to capture both returns and their gradi- ents, coining the framework Distributional Sobolev Reinforcement Learning. At its core is novel Sobolev Bellman operator that bootstraps both return and gradient distributions. By combining gradient-based training with uncertainty modeling, we aim boost policy and value learning. This necessitates generative model that supports differentiation of outputs and their input gradients, hence we introduce Distributional Sobolev Training and detail its implementation. Since most envi- ronments are non-differentiable, transitions and rewards. This enriches SVG (Heess et al. 2015) with a more expressive neural ar- chitecture. Finally, extend previous works on value gradient and introduce the framework of Sobolev Temporal Difference. We provide the first contraction proofs in this scheme. In this context, introduce the maximum-slicet MMD metric practical divergence that induces contraction and is tractable to approximate.
employ conditional VAE (cVAE) (Sohn et al. 2015) to model
BACKGROUND
2.1 NOTATION AND RL OBJECTIVE
We consider Markov Decision Process (MDP) with continuous state and action spaces, S and A, transition kernel P: S x A + P(S), reward law R: S A + P(R), and initial distri- bution 1 E P(S). A deterministic policy no: S + AA induces the 7-discounted occupancy
dhe = (1-7) =t=0 7 LLa((se To, (Silver et al. 2014 D'Oro & Jaskowski 2020). The
Q-function (" a) is the expected future return starting from state s and action a, l.e., Q"
E E ('1( 4x) 80 =, A0 a). It yields the objective
(0) = Ea-E [Q"* (s. Ta Ta()))
(1)
Under mild conditions Silver et al. 2014), the Deterministic Policy Gradient theorem gives 1
VoJ(O) 1 -y E Band Voo V.Q"*(ss, (s)
2.2 TEMPORAL-DIFFERENCE AS AN AFFINE OPERATOR
In practice, the true Q is unknown and approximated by parameterized critic Qg More gener- ally, any value-like mapping V: S x A + y (where y = IR or space of probability distributions Bellemare al. 2017)) admits a temporal-different
update written single affine operator:
(Ta VT, a) = b(s, a) + +[V]ss),))
Here b(s, u) injects the immediate-reward term and L linearly transforms the successor estimate. The following recovers the Bellman expectation operator where Q is the state-action value function
b(s, a) = E. [[{ (lss.)), a)),
[EXP[Q](s a) E[Q(s', क (s')) s,a]
(4)
In distributional RL Bellemare al. 17t (s, a). One gets
2017). let 7 (s, be the random return with distribution
b(s, a) Law [R(s, a)),
CDi [n](ss))) Eg (x + 70) ((ss==(()))
(5)
where (or + 7 is the law of x when X ~ which yields
(T.Dim (s, a) Law [R(s, a) + + Z(s', (0,,,=(())) <s/))], where
8,
(6)
Using off-policy samples (s,a,r,8') B from replay buffer (Mnih et al. 2013] together with
delayed target networks 0', ø' (Lillicrap et al. 2016 Fujimoto et al.) 2018 Haarnoja et al. 2018).
define the one-step targets
881(8, s') r + Qø (8', morr))))))
Th Law [r + 7 (8', no (s)))-
(7)
2

---

Published as a conference paper at ICLR 2026
The critic Va (scalar Qs distribution 2ω) is then trained by minimizing
(()) ) (Vo(s, a), T(s,a)) ((s,a))]
where 0igi the expected-value case or 7 = TAtt in the distributional case, and d is either a regression loss (e.g. squared error) or distributional metric such as the Wasserstein distance Bellemare et al. 2017; Sun al. 2024
A NEW BELLMAN OPERATOR
3.1 LEARNING USEFUL CRITIC
Many value-based methods (Lillicrap et al. 2016: Fujimoto et al. 2018 Haarnoja et al. 2018 rely on learned critic provide the actor's training signal, implying that "an actor can only be as good as allowed by its critic" OD'Oro & Jaskowskij 2020). Unfortunately, typical critics predicting only mean returns - cannot capture inherent return uncertainty. Distributional RL addresses this by modeling the return distribution. However, as noted in D'Oro & Jaskowski (2200) another fun- damental issue that minimizing TD-error does not guarantee
critic will be effective at steering policy optimization. We therefore propose more principled approach distributional temporal- difference learning, which explicitly incorporates the critic's 's action- grrdiinnt into its training objec- tive. This aligns critic optimization with policy improvement rather than just fitting returns their distribution.
Proposition 1. Let π be an L Liipsciit continuous policy, and let Law (V a (8, la=x())) and Law Naw[V..(s, a) la=n(s)] denote the true and estimated distributions of the action-gradient at a = T(S), respectively. Define the p- Wassersteii distance between two probability measures μ., by
Vy(,,v) = ( inf [xx,y~<<lxx (XIIX
YIP!)
1/P
(CElI(a))))
Then, specializing to p 1, the error between the true policy gradient [ and its estimate V (o) satisfies
VoJ voJ(o)| S
L Eundy W (Law V V. (s, (== Law Now[VaZ(ss)) (s, lo=
The proof is in Appendix D This result generalizes Proposition 3.1 D'Oro & Jaskowski (2020) distributional setting. The Lipschitz continuity of typically holds when using network function approximation. We discuss why this assumption is reasonable in practice in Appendix C.2
Following O'Oro & Jaskowski (2020), induce critic optimization objective from Proposition 1 showing that we can approximate the true policy gradient by matching the action gradients in the dis- tributional sense. Using bootstrapping to approximate the true distribution leads the optimization problem
7 E arg min E. Badd Wi (Law (.. 21, a) (= Law [V. Fnge (8, l
ZEz
Zig a) = a(s, + (s', mo(ss)) (s',)) s(.1).a).
(10)
We note that Equations 9 70 assume a known and differentiable dynamics model p. We maintain this assumption for the time being and will relax it in Section Mirroring previous work (D Oro Jaskowski 2020 Garibbo el al. (024), we introduce this assumption upfront and then lift the constraint. In the next section, we formalize the notions necessary to instantiate working imple- mentation this optimization problem.
3.2 DISTRIBUTIONAL SOBOLEV TRAINING
In this section, introduce novel Bellman operator for learning the joint distribution of the discounted cumulative reward and its action getionngraddent, and then express it in the dfffine-ransfor form presented earlier.
3

---

Published as a conference paper at ICLR 2026
Random action Sobolev return We extend the random return [(s, to a joint random variable that captures both the return and its action siriinngradiet.... Formally, the random action Sobolev return
20
(s
Σ r oz):
Σ t=
> r(ss a))
80 = s, do
(11)
Next, we define the Sobolev distributional Bellman operator T over these (I.A] 1) -ddimnssional random variables. Let (s, a) = Law [ (s, a)]- We borrow notation from Zhang et al. (202T); Rowland et al. (2019) and extend the classical distri- butional operator follows. Under policy n, sample
Sobolev distributional temporal difference
~ (.| s, a),
= T(''),
T~ R. s,)), X' ~ Po
Define the full affine transform
5 (2; (,; ,), ( [ faction (x)
(12)
For readability we hereafter write g (x:), implicitly carrying the dependence on (r, 7))). We define T as the operator that, at each (s, a), pushes the next-step law 17 a²) forward through this pointwise affine map:
(T 17 (s.) Law[fS.
(13)
Its components are
preturm (I) r + xx return
(14)
paction (x)
d (8,a) + Da
Ör
Of ( (s, 10 + (8,T(8)) action
(15)
This action-gradient component is novel: it arises by differentiating the Bellman target, capturing how the return's gradient transforms under P(s', a). The derivation Eqs. 14 appears in
Appendix C Notably, since f, and π are differentiable. these updates are implemented automati- cally via backpropagation through the reparameterized simulator (Baydin et al. 2018 Paszke al., 2019 Bradbury et al. 2018).
Affine form of the Sobolev Bellman backup backup can be written in single affine-operator form:
shown in Appendix C the Sobolev Bellman
(s, b(s,a) + (s.
(16)
where b(s, a) (r(s,a), )r)(s,a)) RI+AI collects the immediate reward and its action-
gradient, and (Sob a) is a state-actie ieeee-ccioon--ee deeeeddd linear operator (not merely matrix) encap- sulating the Jacobian blocks of the transition f and policy т. Importantly, Equation 16 represents clear departure from recent works that use critic gradient information Oro & Jaskowski 2020) Garibbo et al. 2024). Whereas these methods use the action- -oriiient only as an auxiliary regulariza- tion signal, our formulation incorporates it directly into the quantity we perform temporal difference over, jointly with the scalar return. This joint TD structure is what enables the contraction analysis presented in the next section.
Complete Sobolev TD. is worth noting that exactly the same chain-ruke derivation used in Eq. 15 extends to also bootstrap the state gradient (not just the action- gradient), yielding the complete Sobolev Bellman operator. We discuss this operator in detail in Appendix C In the main text, we focus on the incomplete version for clarity and computational tractability. Although using both action- and state-gradients would provide more information, handling them together is substantially more computationally expensive and is known to be non-trivial in practice (Garibbo et al.) 2024). The complete operator therefore provided theoretical framework, while the incomplete version captures the core ideas and is sufficient for our algorithmic development.
generalization for readers interested in the full
THEORETICAL RESULTS
The most natural distance for distributional RL is the Wasserstein metric, and therefore begin by establishing contraction under this choice. We state the contraction results here, and defer the ex- plicit contraction factors and the smoothness assumptions they rely on Appendix The appendix
4

---

Published as a conference paper at ICLR 2026
also contains results of the same type for the complete operator. The main takeaway is that, under appropriate and practically reasonable assumptions, we can show contraction.
Theorem (Action-gradiens Sobolev contraction). Assume bounded Jacobians for the transition f and policy n, and Lipschitz coupling relating state- to return- Sturn-gradients, Then
Wo (I m. T2"m) S 7/K W (2))
where W ( 72) (a.a) (n a), 12(1, a)) and depends the smoothness assumptions. Y 1, T a strict contraction.
Towards tractable metric Both the classical and Sobolev distributional Bellman operators (Eqs. 6 [3) involve intractable push-forward integrals. While the supremun-p-Wasersten dis- tance W, natural theoretical choice (Bellemare al. 2017) exact multivariate optimal trans- port costs O(m² logm) with samples and, more broadly, Wasserstein distances difficult to estimate and to use directly for training (Appendix E). This motivates a shift to metrics that re- main faithful distributional structure yet easier to compute in practice. One such candidate is the Maximum Mean Discrepancy (MMD) (Gretton et al. (0122). kernel-based divergence that tractable, sample-based and already explored in distributional RL (Nguyen et al., 2020 Killing- berg Langseth 2023; Wiltzer et al. 2024). For laws P.Q C IR"", the squared Maximum Mean Discrepancy between P and Qis
MMD Q) E,,"' k(xx x')] + kk v²)] 2 E.- v- w)],
(17)
where & denotes the kernel function. Further properties and empirical estimators are discussed in Appendix F2
Max-Sliee MMD.
To obtain provable contraction for Sobolev operator, lift MMD via
the max- slx-dliced divergence framework (Deshpande al. 2019 Nadjahi al. 2020). For 0 E g and (x) = (0, z). we define
MSM = sup MMD( (Po) (P))
Approximation. MSMMD be approximated by gradient-based optimization of the direction on the unit sphere; see Algorithm
Contraction under Max-Sliced MMD.
Theorem 2 Daction-gradien Sobolev contraction under MSMMD), Assume the conditions of The- orem I and the mild additions in Theorem 6 Then
MSMMD T m. T Y MSMMD m7 n2).
where MSMMD (m, 72) = sup (aa)ESXA MSMMD(n (s, (s,a), 12(8, a)) and n. depends on the Ja- If T is strict contraction with unique fixed point.
cobian bounds; see Appendix G
Trade-off interpretation The condition к < 1 makes the trade-off explicit: either enforce smoothness (reduce к. via bounded Jacobians and Lipschitz couplings) or shorten the effective hori- zon (reduce 7). A similar quantity к governs our contraction results for both Wasserstein (Theo- rem [T and max-sliced MMD (Theorem 2 so the trade-off arises in each case. Importantly, x is dictated by the environment's dynamics and policy sensitivities; when the underlying physics yields large gradients, this effect cannot be eliminated and the only remedy is to lower 7. This observation a contribution of our work and is enabled by the Sobolev Temporal Difference framework introduce.
5 APPROACH
To turn our theoretical Sobolev-distribttions Bellman operator into a practical algorithm, we need (i) critic that produces joint return-gradient samples, (ii) Bellman backups without intractable
5

---

Published as a conference paper at ICLR 2026
integrals, (iii) a tractable metric for comparing predicted and target distributions, and (iv) to relax the assumption that our environment is differentiable. We address each of these below.
Sobolev inductive bias. We coin the term Sobolev inductive bias for the simple idea of modeling a gradient by a gradient: using the gradient of our approximator to stand in for the gradient of the true function. Concretely, let F: IR - R6 be differentiable target and fP neural network parameterized by s. Sobolev training is one concrete instance of this principle (e.g. (Czarnecki
2017 D'Oro & Jaskowski 2020)), which implements (p;;)) - [I (x))] As F
al.
V fff (x))), Further mathematical motivation for this inductive bias can be found
in Appendix B.1
Reparameterized Sobolev critic. We model the joint return-and-gradient law n - a) by gen- erator
25. (8, (5 - (Z), V (,
& {~N(0,1)
T (s. This purely sample-base based critic sidesteps intractable likelihoods and scales naturally high-dimenniooal actions. is similar in spirit to Singh it al. (2022); Freirich et al.] (2019). and is structured as generative model that deterministically maps noise to samples (Li et al. 2015, Goodfellow et al., 2020). We discuss this choice further and why most alternative parametrizations would not fit our scenario in Appendix B.2
such that Law [ (s. a)] = (..(s.))
Overestimation bias. Value estimates routinely exhibit overestimation (Hasselt) 2010 van Hasselt et al. 2015), and even grgullnt-regularrzed critics inherit this issue (D'Oro & Jaskowski 2020] Garibbo et al. 2024). biases the policy toward overvalued actions, undermining its performance. TD3 addresses by training two critics and setting the target in Equation ( to the minimum of the two (Fujimoto et al., 2018). In our sample-based distributional setting, we follow TQC (Kuznetsov) ct al. 2020 ; train two distributional critics, draw N samples from each, discard the top p%o by magnitude, and concatenate the rest to form the target distribution.
One-step world model. Now we relax the assumption that the environment is differentiable. Un- like (D'Oro & Jaskowskil 2020 Garibbo et al., 2024). which fit only the conditional expectation f) Blg"rr r 8, a). we learn stochastic, differentiable simulator g whose push-forward law approximates the true transition-reward distribution P(s', ,r a. Concretely, we posit
() = g(s, a, E), E~ (E).
Law [g(s, a)] = Law (s'rr a].
Using g in place of the true environment, our Sobolev Bellman update from Equation 12 becomes
nde (s. a) = So (Z (8', d'); f, 7),, à' = (8').
We implement 9 via a conditional VAE Sohn et al. 2015).
which has proven effective for modeling
dynamics in RL (Ha& & Schmidhuber] 201 8 Zhu et al. 2224). with learned prior (o(e s, a.), encoder
gc(E| s, s, s', r), and decoder (s",r | s,a,e).
Because Sobolev TD requires drawing many model transitions per update and differentiating each sample with respect to (s, a). the world model must support both cheap sampling and cheap repa- rameterized gradients. Several generative model classes satisfy these requirements, such as normal- izing flows (Dinh et al. 2016) or GANs (Goodfellow et al. (020)). Diffusion models (Ho et al. 2020) have shown strong empirical performance in high-dimensional generative tasks, but they are not well suited to our setting: computing the Jacobian of generated sample with respect to its con- ditioning variables requires backpropagating through the full denoising chain, which is prohibitively expensive for the repeated one-step model queries whose gradients Sobolev TD evaluates.
Algorithm summary. We replace the critic in standard off off-policy actor-critic (Fujimoto et al. 2018 Barth-Maron et al. 2018) with our Sobolev distributional critic and train it and the VAE world model gs, a, E) jointly; this yields the Distributional Sobolev Deterministic Policy Gradi- ent (DSDPG) algorithm. At each step we sample Sobolev-returns and (8', f), form bootstrapped Sabolev-Bellman targets, minimize the MSMMD loss on salue-gradient pairs, and update the actor by ascending the gradient of the critic's estimated expected return. Even though plain MMD is not shown to be contractive, we also instantiate our framework using it, thus proposing two metrics to train the critics. Fig. [ (left) illustrates the overall DSDPG workflow, while the right panel shows how Sobolev-retume retuur distributions are sampled and the MMD loss is estimated. The full pode can be found in Appendix |

---

Published as a conference paper at ICLR 2026
smeple predtcted
tgt,, peltoy. scrplex froe
(erstie
Dermpency
cooputa paltey(k2)
Pany 19 (
target zcrples
Jsel
crstics_
froe cangeta
gaees
per sxepse
autograd grad
joint
World (eVAR)
prone topupl
concatenata
prame.top.peee top.percent
Figure I: Left: block diagram of our DSDPG algorithm, where the critic Zo maps noise & and (s, a) to Sobolev-return samples, vAE world model generates next-staterreward samples (8', MSMMD MMD compares predicted and bootstrapped Sobolev-return distributions, and the pol- icy is updated via the critic's mean (gradient flows shown as dashed arrows; inspired by (iinhh et al. 2022 ) Right: pseudocode for estimating the biased MSMMD between predicted and bootstrapped
Sobolev returns.
RESULTS
6.1 TOY REINFORCEMENT LEARNING
We introduce simple continuous- state, continuous-action 2D point-mass task within a square bounding box (Figure Zap. At each time step, the agent controls the 2D acceleration of the mass. At the start of each episode, we randomly draw one of N possible bonus locations arranged uniformly around the center. Reaching this location ends the episode with terminal reward. The agent does not know which location is correct. Instead, each time it visits location, binary memory flag set. This partial observability makes the task sparse and forces exploration. By varying N, we tune the number of modes in the return distribution: small N yields near-deterministic returns, large N yields multimodal, high-variance returns. This toy task offers clean, adjustable benchmark for gradient-aware distributional critics under controlled multimodal uncertainty.
As shown in Figure 2b) Distributional Sobolev (MSMMD Sobolev in pink and MMD Sobolev, or- ange) consistently outperforms all baselines we increase the number of bonus locations. This indicates that our method effectively leverages gradient information and remains robust to the grow- ing multimodality red) offers no clear advantage over gradient-agnostic critics. Interestingly, the provably contractive MSMMD has mild advantage over plain MMD. For brevity, defer the full experimental details and baseline descriptions to the next subsection, while exact empirical settings are in Appendix K
the return distribution. By contrast. deterministic Sobolev (Huber Sobolev,
N-3
H-4
N-5
N-6
Colicion Momory 0.
banus
1911
Muaeet trtneer PHAGO
(8)
(b)
Figure 2: (a) Illustration of the 2D point-mass environment with N possible bonus locations. (b)
Evaluation curves (10 means the agent reached the bonus) for our Distributional Sobolev (MMD
Sobolev and MSMMD Sobolev, in orange and pink resp.), deterministic Sobolev (MAGE Oro & Jaskowski 2020), in red), and other basclines N varies (median over 5 seeds and 25%-75% IQR).

---

Published as a conference paper at ICLR 2026
6.2 REINFORCEMENT LEARNING
In this section, we evaluate the complete solution, including the learned world model, on several standard MuJoCo environments (Todorov et al. (2012) from the Gymnasium library Towers et al. (2024). Our deep learning framework is JAX (Bradbury et al., 2018). Similarly to MAGE (D'Ore & Jaskowski, 2020p, for every exploration step. 10 critic updates are taken (UTD ratio 10). All experiments are run in a Dyna setting @Deisenrthh & Rasmussen 2011) where new samples are drawn from the world model for every critic update. Thus, we evaluate the ability of the Distri- butional Sobolev training framework against deterministic Sobolev (Czarnecki et al. [017]] D'rro & Jaskowski, 2020, Garibbo et al., 2024). Hence, the results should not be understood as looking for ofite-of-the-ar performance in the given environments but rather a showcase that in some difficult settings, distributional methods can perform better.
We compare against four basclines: (i) TD3 variant trained with a Huber loss (Fujimoto et al. 2018 D'Oro & Jaskowski] 2020); (ii) deterministic Sobolev (MAGE) trained on both return and action-gradient via Huber loss (D'Oro & Jaskowski| 2020); and (iii-iv) distributional baselines, IQN (Dabney et al. 2018a) and standard MMD training (Nguyen et al., (020). All variants involving MMD use multiquadric kernel with h 100. We use identical base architecture for every method, with the distributional approaches sampling by reparameterization (Section 5 Each dis- tributional variant generates 10 samples per transition. Further architectural and hyperparameter details are in Appendix L
To evaluate robustness under stochastic dynamics, we use two complementary perturbations in Mu- JoCo. (i) Multiplicative observation noise: at the start of each episode, we sample n 2[[0.8, 1.2] and scale the observation s + n s, following Khraishi & Okhrati (2023); this induces partial observ- ability and substantially increases task difficulty. (ii) Additive Gaussian dynamics noise: we inject zero-mean Gaussian noise directly into positions and velocities, adopting the standard deviations reported by Khraishi & Okhrati (002)). This second perturbation makes control harder while pre- serving the qualitative structure of the original tasks. Results are summarized in Figure 3 where the left panels report the final evaluation performance after 150,000 iterations (250,000 for Humanoid- v2), and the right panels show the normalized area under the evaluation curve (nAUC). All results reported medians with bootstrapped 95% confidence intervals. The full learning curves provided in Figure T0.
In the noise-free setting, DSDPG (MSMMD Sobolev and MMD Sobolev) match the performance of all baselines across six MuJoCo tasks. Under multiplicative observation noise, it outperforms every competitor in three of the six environments, most notably Ant-v2 and Humanoid-v2, while deterministic Sobolev (MAGE) suffers severe drops on Walker2d-v2 and Humanoid-v2 and shows Under Gaussian noise, DSDPG again
larger variance in the simple InverieiDoubbeeendllum-2 outperforms the baselines in three of the six environments. Notably, on Ant-v2, the noise scale borrowed from Khraishi & Okhrati 2023) makes the task substantially more difficult than previously reported.
Additional experiments are provided in Appendix M including sensitivity to the noise scale, kernel bandwidth, number of Sobolev samples, and world-model capacity. We also report an ablation that removes the overestimation bias correction. For DSDPG, this corresponds to disabling the TQC truncation (Kuznetsov et al. 2020), and for deterministic baselines, to removing the double estima- tion (Fujimoto et al. 2018). This ablation shows that overestimation bias correction plays a crucial role in stabilizing learning and achieving high performance. A wall-clock runtime comparison for Humanoid-v2 is provided in Table 3]
World-model ablation. To verify that our method does not depend on the particular choice of <VAE as world model, we performed an ablation in which the c cVAE was replaced with a lightweight normalizing- flow model (Dinh et al. 2016). The flow-based model integrates into the Sobolev TD framework without modification, as the algorithm only requires that the world model provide cheap sampling and differentiable one-step transitions with respect to (s, a). Across all tested MuJoCo tasks, the qualitative behaviour of the method remains unchanged, indicating that the benefits of Distributional Soboley TD arise from its ability to exploit gradient information in stochastic envi- ronments rather than from the specific generative architecture used. A brief presentation of normal- izing flows is provided in Appendix(H22 [H.2], and the full ablation results, including architecture, training setup, and figures, are reported in Appendix M.
8

---

Published as a conference paper at ICLR 2026
Maturaaln2
5000
2000
Hornal
4000
2002
Merrral
Mallte cartie no
CoLin or notie
Hepper-tt
Hepperout
2003
Nomo
Haltchettahra
MAIC
6000
Marrrol
Milligg citian
Gaustar note
5903
Hemmal
Mul Poctti
Multte
Aetiv
Ant
1000
Mermal
Mcliill catliee
MOHMD
(MAGF)
IQN
MSNHD
Sebolie (MAGE])
NISNND Scbelev
HND
Huber 1TD30
NVD Sobolcy
NVD
HUSOr ITD3)
VHD Sobslee
(b)
Figure 3: Evaluation of DSDPG (MSMMD/MDD Sobolev), deterministic SobolevMMGE D'Oro & Jaskowski (2020). 393-HuberFujimotte et al. (2018). ION et al. (2018a), standard MMD Nguyen et al. (2020); Killingherg & Langseth (2023) six MuloCo tasks. Results reported over 10 random seeds with median and 95% bootstrap confidence intervals. We compare three settings: normal environment, multiplicative observation noise, and Gaussian dynamics noise (Luo ct al. 2021). (a) Final evaluation performance. (b) Normalized AUC over the entire training curve. Our DSDPG variants (MMD Sobolev and MSMMD Sobolev, red and brown) are on par or better than competing methods, shining especially on harder tasks and under noisy environments.
7 RELATED WORKS
Gradient informed training Our work explicitly models gradients a stochastic manner using neural networks. As such, it can be
distributional extension of Sobolev training Czar-
necki et al. (2017) which was already adapted to reinforcement learning in D'Orro & Jaskowski
2020)]/arribb Garibbo al. (2024). Gradient modeling in value-based RI. dates back to Fairbank (2008). Additionally, our approach shares connections with Physics Informed Neural Networks (PINNs)
RRaissi al. 2017). which approximate physical processes via differential constraints. In particu- lar, uncertainty-awaat PINNs (Yang et al.) 2020 Zhong & Meidani] 2023) treat random processes and their derivatives random variables, akin to our approach of leveraging Sobolev inductive bias and generative modeling (through reparameterization) while enforcing consistency via tractable distributional distance.
Distributional RL We extend distributional RL (Bellemare et al. 2017 2023 by modeling return gradients for deterministic policies continuous action spaces, building on distributional DDPG (Barth-Maron et al.] [0188] [illicrap al.) 2016). Because these gradients are multi multi-dimmesional. our method aligns with distributional multivariate returns (Zhang et al.) 2021 Freirich et al.) 2019 Wiltzer et al. 2024), highlighting the need for tractable discrepancy measures over multi- dimensional distributions. We treat our distributional critic as a generative model capable of produc- ing actual samples of the underlying distribution (Freirich et al. 2019; Singh al. [022] Doan et al. 2018). in contrast to approaches that generate only pseudo-sane samples (Zhang et 20.21 Nguyen al. 2020) rely summary statisties (Bellemare al. 2017 Barth-Mann Maron al. 2018 Dab-

---

Published as a conference paper at ICLR 2026
ney al., 2018b a). We measure distributional discrepancy via the Maximum Mean Discrepancy (MMD) Grrttoon et al. 2012 Li et al. 2015; Binkowski et al. 2021; Oskarsson 2020). which has
proven effective in distributional RL Nguyen et al. 2020: Killingberg & Langseth 2023 Zhang
et al.
2021
Model-based RL Finally, because environment dynamies and rewards unknown and non- differentiable, we adopt a world model akin to SVG(1) (Heess et al.] (015), instantiated as a cVEE Kingma 2013 Sohn et al. 2015). Generating new data through this model places our approach within model-based RL (Chua et al. 2018 Feinberg et al. 2018). particularly the Dyna family (Sut-) ton 199I). More specifically. our method related to approaches that rely on variational techniques
(Ha & Schmidhuber 2018 Hafner et al. [2020; Zhu et al., (2024) or backpropagate through world
models (Hafner et al. 2019 Clavera et al. 2020: Amos et 2021 Henaff et al. 2017 2019
Byravan et al. 2020).
8 CONCLUSION, LIMITATIONS AND FUTURE WORK
In this work, we introduced Distributional Sobolev Deterministic Policy Gradient (DSDPG). Our main contributions involve modeling a distribution over both the output and the gradient of ran- dom function, and deriving a tractable computational scheme to achieve this using Maximum Mean Discrepancy (MMD) and its max-sliced variant MSMMD. We demonstrated the effectiveness of the method in leveraging gradient information. We further extended the approach to reinforcement learning by leveraging differentiable world model of the environment that infers gradients from observations. Building on this, we proposed the framework of Sobolev Temporal Difference, allow- ing us provide the first contraction results in gradient-aware RL. We established contraction for both Wasserstein and the tractable MSMMD divergences, and highlighted smoothness assumptions required for contractive gradient-aware training. Finally, we showed that distributional gradient modeling improves stability in a controlled toy problem and enhances robustness to noise in Mu- JoCo environments, with stronger benefits in high-dimensional tasks.
While our approach shows promise, there remain challenges to address. A primary consideration is the high computational cost, as both policy evaluation and improvement require multiple samples from the distributional critic and their input gradients. Future work should explore more efficient inductive biases.
A second avenue for future work concerns the different variants of the Sobolev Temporal Difference operator. In the main text, we focus on the action-gradien variant for clarity and tractability, whereas the complete variant also includes state-gradics giiedgrraients and provides additional information about how returns vary with the dynamics. Using both gradients jointly is computationally demanding and remains non-trivial in practice, but addressing these challenges could further extend the applicability of Sobolev-based distributional critics.
Finally, we believe that the ideas introduced in this work could benefit other fields where aleatoric uncertainty in gradient modeling is important, such as Physics-Informed Neural Networks ( (Yang & Perdikaris 2019), Neural Volume Rendering (Lindell et al. (021) and more generally applications of double backpropagation (Drucker & Le Cun 1991).
ACKNOWLEDGEMENTS
This project has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement n² 101021347).
REFERENCES
Charalambos D. Aliprantis and Kim C. Border. Infinite Dimensional Analysis: Hitchhiker's Guide. Springer Verlag, Berlin, Heidelberg, 3 edition, 2006. ISBN 978-3-540-299586.1 doi: 10.1007/
3-540-29587-9, URL httt:t:::4 /xx wooi.oogg/1.1110
540 29587-9
Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Learning for Dynamics and Control, pp. 6-20. PMLR, 2021.
10

---

Published as a conference paper at ICLR 2026
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Interaational conference on machine learning, pp. 214-222 PMLR, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprini arXiv: 1607.06450, 2016.
Philip J Ball and Stephen J Roberts, Offeon3: What is state of the art anyway? arXiv preprint
MIXiv:2/01.1131,, 2021.
Gabriel
Matthew W. Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva TB, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional deterministic policy gradients. In International Conference on Learning Representations (ICLR), Vancouver, Canada, 2018.
Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey, Journal of Machine Learning Research, 18(153):1-44 43, 2018. URL http:////...l..oooaaaaseee//////
ht.
Mare G. Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforce- ment learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pp. 449 458. PMLR, 06-11 Aug 2017. URL / mir.
be lemare 7a.htm) htm
Mare G Bellemare, Will Dabney, and Mark Rowland. Distributional reinforcement learning. MIT Press, 2023.
Mikolaj Binkowski, Danica J. Sutherland, Michael Arbel, and Arthur Gretton. Demystifying mmd gans, 2021. URL
www.ii.wwrrrroo.. wwws///ll111 01401
Charles Blundell, Julien Comebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1613-1622, Lille, France, 07-09 Jul 2015. PMLR. URL html
mll..
ittess/33//buneee
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations Python+NumPy programs, 2018. URL http:
//github. com/ jax-mm/jaxx
Arunkumar Byravan, Jost Tobias Springenberg, Abbas Abdolmaleki, Roland Hafner, Michael Ne- unert, Thomas Lampe, Noah Siegel, Nicolas Heess, and Martin Riedmiller. Imagined value gra- dients: Model-basee policy optimization with tranferable latent dynamics models. In Conference on Robot Learming. pp. 566-589, PMLR, 2020.
Shreyas Chaudhari, Srinivasa Pranav, and José MF Moura. Gradient networks. arXiv preprint arXiv:2404.03361, 2024.
Xinyue Chen, Che Wang, Zijian Zhou, and Keith Ross. Randomized ensembled double lelarning: Learning fast without a model. arXiv preprint arXiv:2101.05982, 2021.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforce- ment learning in handful of trials using probabilistic dynamics models. In S. Ben- gio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neurul Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https::/pproceeddniit nerrisseerrssss cate..cc/peee wwtpppppppppppeeeeeeeee httpperrfiileespppee//01111
File/ e/ aa22aee6885555911hooe13111210003
Ignasi Clavera, Violet Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. arXiv preprint 2xXi:2005500863 2020.

---

Published as a conference paper at ICLR 2026
Wojeiech M. Czarnecki, Simon Osindero, Max Jaderberg, Gizegorz Swirszez, and Raz- van Pascanu. Sobolev training for neural networks. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https::/ ngs www.:nnwwrrrccceeeo cc/paper.ppppp.pe http://aarrfffsspeaeeee files/paper/20117 FL1//9BBu61163888acce1166e222222
pdf
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In International conference on machine learning, pp. 1096- 1105. PMLR, 2018a.
Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings the AAAI conference on artificial intelligence, volume 32, 2018b.
Mare Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach policy search. In Proceedings of the 28th International Conference on machine learning (ICML-11), pp. 465-472, 2011.
Ishan Deshpande, Yuan-Time Ting Hu, Ruoyu Sun, Ayis Pyrros, Nasir Siddiqui, Sanmi Koyejo, Zhizhen Zhao, David Forsyth, and Alexander G Schwing. Max-sliced wasserstein distance and its for gans. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10648-10656, 2019.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv::605.. 165.00000 2016.
Thang Doan, Bogdan Mazoure, and Clare Lyle. Gan q-learning, arXiv preprint 1805:1800004874 2018.
Pierluca D'Oro and Wojciech Jaskowski. How to learn a useful critic? model-based action-graden gradient- estimator policy optimization. In 34th Conference on Neural Information Processing Systems (NeurIPS 2020), Vancouver, Canada, 2020.
Harris Drucker and Yann Le Cun. Double backpropagation increasing generalization performance. In IJCNN-91-Seatle International Joint Conference on Neural Networks, volume 2. pp. 145-150. IEEE, 1991.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018.
Michael Fairbank. 2008.
Reinforcement learning by value gradients. arXiv preprint arXiv:0803.3539,
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint arXiv: 2018.
Dror Freirich, Ron Meir, and Aviv Tamar. Distributional multivariate policy evaluation and explo- ration with the bellman gan. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pp. 2361-2370, Long Beach, California, 2019. PMLR. *Equal contribu- tion.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference on machine learning, pp. 1587-1596. PMLR, 2018.
Michele Garibbo, Maxime Robeyns, and Laurence Aitchison. Taylor td-learning. Advances in
Neural Information Processing Systems, 36, 2024.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complex- ity of sinkhorn divergences. In The 22nd international conference on artificial intelligence and statistics, pp. 1574-1583. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2026
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-144, 2020.
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-773, 2012. URL
atttp://mmie
html
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Im- proved training of wasserstein gans. Advances in neural information processing systems, 30,
2017.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122. 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learming, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1rXiv:1912.01660 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02 1933, 2020.
Hado Hasselt. Double q-learning. In J. Lafferty, C. Williams, J. Shawe-Taylor R. Zemel, and
A. Culotta (eds.), Advances Neural Information Processing Systems, volume 23. Curran Associates, Inc., 2010. URL bttps::/proccedinggg www.::p.....iiiis.occomoom cc/pper files/ capeer//000FF11eeoooorSreeeeeee2
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learn- ing continuous control policies by stochastic value gradients. Advances in neural information processing systems, 28, 2015.
Mikael Henaff, William F Whitney, and Yann LeCun. Model-based planning with discrete and continuous actions. arXiv preprint arXiv:/705.07177 2017.
Mikael Henaff, Alfredo Canziani, and Yann LeCun. Model Model-predictive policy learning with uncer- tainty regularization for driving in dense traffic. arXiv preprint arXiv: 022005, 2019.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning, 2017. URL
http//aaivv... http://xwxwiorrx//////// httpa/bbs//11..//// 0 02298
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 0166840-6511 2020.
Max Jaderberg, Wojciech Marian Czarnecki, Simon Osindero, Oriol Vinyals, Alex Graves, David Silver, and Koray Kavukcuoglu. Decoupled neural interfaces using synthetic gradients. In Inter- national conference on machine learning, pp. 1627-1635. PMLR, 2017.
Raad Khraishi and Ramin Okhrati. Simple nois enviromment augmetttti for reinforcement learn- ing, 2023. URL htt.:/ss//aar2....o////////22..2 wttp://rrrg/bbb//....22222
Ludvig Killingberg and Helge Langseth. The multiquadric kemel for moment- mattentn distribu- tional reinforcement learning. Transactions Machine Learning Research, 2023.
Diederik Kingma. Auto-encoding variational bayes. arXiv preprint arXiv: 1312.61 2013.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions. Advances in neural information processing systems, 31, 2018
13

---

Published as a conference paper at ICLR 2026
Anders Krogh and John Hertz. A simple weight decay can improve generalization. In J. Moody, S. Hanson, and R.P. Lippmann (eds.), Advances in Neural Information Processing Systems, vol- ume 4. Morgan-Kaufmann, Kaufmann. 1991. URL
files/ httpaa//SS//.eeeedcdooooooooeee0
Arsenii Kuznetsov, Pavel Shvechikov, Alexander Grishin, and Dmitry Vetrov. Controlling over- estimation bias with truncated mixture of continuous distributional quantile critics, 2020. URL
httpp:://ariv..oe wttpa//wwwwww.r.a///////oooaa a0tp:////.2200..... 04269
Yujia Li, Kevin Swersky, and Rich Zemel. Generative moment matching networks. In International conference on machine learning, pp. 1718-1727, PMLR, 2015.
Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa,
David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Inter- national Conference Learning Representations (ICLR), San Juan, Puerto Rico, 2016.
David B Lindell, Julien NP Martel, and Gordon Wetzstein. Autoint: Automatic integration for fast neural volume rendering. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 14556-14565, 2021.
Liyuan Liu, Haoming Jiang. Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. arXiv preprint arXiv:1908.00 03266, 2019.
Yudong Luo, Guiliang Liu, Haonan Duan, Oliver Schulte, and Pascal Poupart. Distributional rein- forcement learning with monotonic splines. In International Conference on Learning Represen- tations, 2021.
Anton Mallasto, Guido Montúfar, and Augusto Gerolin. How well do wgans estimate the wasserstein metric? arXiv preprint w:Xiv:1910003875 2019.
Takeru Miyato, Daisuke Okanohara, Shin-ichi Maeda, and Masanori Koyama, Synthetic gradient methods with virtual barward-backward networks. 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka. Nonparametric return distribution approximation for reinforcement learning. In Proceedings of the 27th International Conference on Machine Learning (ICML-10), pp. 799-806, 2010.
Kimia Nadjahi, Alain Durmus, Lénaïc Chizat, Soheil Kolouri, Shahin Shahrampour, and Umut Simsekli. Statistical and topological properties of sliced probability divergences. Advances in Neural Information Processing Systems, 33:20802-20814 2020.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference on machine learning (ICML--00, pp. 807-814, 2010.
Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning with maximum mean discrepancy. Association for the Advancement of Artificial Intelligence (AAAI), 2020.
Joel Oskarsson. Probabilistic regression using conditional generative adversarial networks, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Maziar Raissi, Paris Perdikaris, and George Em Kamiadakis. Physics informed deep leam- ing (part i): Data-driven solutions of nonlinear partial differential equations. arXiv preprint arXiv: 171 1r1iv:1711. 10561, 2017.
14

---

Published as a conference paper at ICLR 2026
Prajit Ramachandran, Barret Zoph, and Quoe V Le. Searching for activation functions. arXiv preprint arXiv: 1rX.v:1710.05941, 2017.
Mark Rowland, Robert Dadashi, Saurabh Kumar, Rémi Munos, Marc G Bellemare, and Will Dab- ney. Statistics and samples distributional reinforcement learning. International Conference Machine Learning, pp. 5528-5536. PMLR, 2019.
René L Schilling, Renming Song, and Zoran Vondracek. Bernstein functions: theory and applica- fionx, volume 37. Walter de Gruyter, 2012.
Dino Sejdinovic, Bharath Sriperumbudur, Arthur Gretton, and Kenji Fukumizu. Equivalence of distance-based and rkhs-based statístics in hypothesis testing. The annals of statistics, pp. 2263- 2291,2013.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pp- 387-395. Pmlr, 2014.
Rahul Singh, Keuntaek Lee, and Yongxin Chen. Sample-based distributional policy gradient. In Learning for Dynamics and Control Conference, pp. 676-688. PMLR, 2022.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances neural information processing systems, 28, 20155
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 19656):1929-1958, 2014. URL http::/jmiz.crr wttp://ppppp/s.///////
srivastaval4a.hne
Jan Stanezuk, Christian Etmann, Lisa Maria Kreusser, and Carola-Bibiane Schönlieb. Wasser- stein gans work because they fail (to approximate the wasserstein distance). arXiv preprint arXiv:2103.01678, 2021.
Ke Sun, Yingnan Zhao, Wulong Liu, Bei Jiang, and Linglong Kong. Distributional reinforcement learning with regularized wasserstein loss. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Richard Sutton. Learning to predict by the methods of temporal differences. Machine learning. 3:9 44, 1988.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160-163, 1991.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5036.5033. IEEE, 2012. doi: 0009//RRS.201226386611
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: standard interface for reinforeement learning environments. arXiv preprint arXiv:2407.17032 2024.
Hado van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. CoRR, abs/1500.06461, 2015. URL http:://aaxxi..oo http://wwww....aa/////a////
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-
learning. In Proceedings of the AAAI conference artificial intelligence, volume 30, 2016.
Cédric Villani et al. Optimal transport: old and new, volume 338. Springer, 2009.
Harley Wiltzer, Jesse Farebrother, Arthur Gretton, and Mark Rowland. Foundations of multivari- ate distributional reinforcement learning, 2024. URL https: axxiv. wrtps///2rrabb/2400...
00328
15

---

Published as a conference paper at ICLR 2026
Liu Yang, Dongkun Zhang, and George Em Karniadakis. Physics-informed generative adversarial networks for stochastic differential equations. SIAM Journal on Scientific Computing, 42(1): A292-A317, 2020.
Yibo Yang and Paris Perdikaris. Adversarial uncertainty quantification in physics-informed neural networks. Journal of Computational Physics, 394:136-152, 2019.
Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Yan Liu. Distributional re- inforcement learning for
reward functions. Advances in Neural Information Processing Systems, 34:1519-1529, 2021.
Ruijie Zheng, Xiyao Wang, Huazhe Xu, and Furong Huang. Is model ensemble necessary? model-based via a single model with lipschitz regularized value function. arXiv preprint
arXiv:2302.01224, 2023.
Weiheng Zhong and Hadi Meidani. Pi-vae: thysics-informed variational auto-encoder for stochastic differential equations. Computer Methods in Applied Mechanics and Engineering, 403: 15664, 2023.
Ting Zhu, Ruibin Ren, Yukai Li, and Wenbin Liu. model-based reinforcement learning method with conditional variational auto-encoder. Journal of Data Science and Intelligent Systems, 2024.
16

---

Published as a conference paper at ICLR 2026
APPENDIX CONTENTS
Derivatives of reparameterizable random variables
18
Sobolev Inductive Bias and Distribution Parametrization
19
B. Sobolev inductive bias
19
B.2 Distribution parametrization
19
Sobolev Bellman operator
20
C.1
Preamble and assumptions
20
Lipschitzness of neural ii..............................
and critics
20
Derivation of the action-gradient term
21
C,4
Distributional Sobolev Bellman update
22
C.5
Distributional Sobolev Bellman operator
23
C.6
Complete Distributional Sobolev Bellman operator
23
D Results with Wasserstein
24
D.I Proofs of contraction and fixed point
24
D.2 Proof of Proposition 3.1
32
Practical difficulties with Wasserstein for training
35
Background on MMD
36
Contraction under MMD
37
Empirical estimators MMD
38
Results with Max Sliced M.D (MSMMD)
39
H
Background on the world-model
47
Conditional variational auto-encoders
47
H.2
Conditional normalizing flows as world models
47
Algorithm design, pseudo-code and baseline
48
I.1
Design choices
48
Pseudo-codes
49
Baseline- MAGE
51
J Toy supervised learning
52
Common tricks
53
J.2 Adding noise
54
Toy Reinforcement Learning
56
K. Environment
56
17

---

Published as a conference paper at ICLR 2026
K.2 Experimental details
56
Reinforcement learning experiments
56
Full curves and wall-clock time
57
M
Additional experiments RL
57
M.I
Conditional flow world model
M.2
Overestimation bias
60
M.3
Multiquadrie kernel bandwidth
60
M.4
Noise scale
60
M.5
Number of samples
63
M.6
World model capacity
63
LLM usage
63
A
DERIVATIVES OF REPARAMETERIZABLE RANDOM VARIABLES
In several places, such as Section we need to talk about "the gradient of a random variable with respect its conditioning variable." To make this precise and intuitive, we rely on the common repa- rameterization trick. We explain that, under appropriate continuity and measurability assumptions, this construction yields a well-defined random variable.
Reparameterization view. Suppose Y x is conditional random variable, where X' C IR" and y/ C R We assume there exists measurable "noise" variable E ~ p(E) on space Z and deterministic map
g x -
that is C in its first argument, such that
0(x,E) has the same law as In other words, for each sample &, the realization
v((:g) ((rr
a deterministic function of
Pathwise derivative. Because g( is differentiable in x. we define the pathwise reparameteri- zation derivative by
i Y def 2
Or g1,
dr
that is, for each fixed 5, O2y(2; {) = (z E). We differentiate the pathwise realization æ H g(x, E), then regard 0_g(2, E) random variable taking values in R
Measurability. Because C in its first argument, for each fixed & the map
T
0,g(x,E)
continuous in x. Moreover, by assumption & + ( {) is measurable for each æ. It follows from Lemma 4.51 from Aliprantis & Border 2006 ("Carathéodory functions are jointly measurable") that
(x, d,
is jointly menelleeas X' x Z. Consequently, when =\t p(E) the composition 0,g(2, E) is a well-defined pp random variable.
Practical computation. In modern autodiff frameworks, once we formulate the generative model Y {(x,)), calling grad(Y, automatically returns the sample-wisg wise Jacobian 0zg(2, { Hence, all "random variable gradients" in our Sobolev Bellman operator are tractable to sample using standard reparameterization and automatic differentiation techniques.
18

---

Published as a conference paper at ICLR 2026
B SOBOLEV INDUCTIVE BIAS AND DISTRIBUTION PARAMETRIZATION
B.1 SOBOLEV INDUCTIVE BIAS
We discussed the Sobolev inductive bias in Section 5 Recall that given differentiable target F: Ra - RR, Sobolev training (Czarnecki et al. 2017) (originally introduced as "double back- propagation" by Drucker & Le (191(19 minimizes both value- and gradient-missatee
(p;;)) = [1)()) fx()) + \s Vafp(z)11.
The first term aligns network outputs to F(x), while the second enforces that the input-dervaative of fo match those of F. Below, we build on this bias to derive conditions on valid gradient fields and their relation Jacobian symmetry.
Using the gradient of neural network to predict gradients is not the only way to do so. Indeed, Jaderberg et al. (2017) proposed generate synthetic gradients directly via the outputs of neural network single linear layer. However, Czamecki al. 22;1);; MMyato et al. (017)) observed that predicting synthetic gradients as gradients of a neural network directly improved performance on decoupled trainings experiments from Jaderberg et al. (2017). A partial explanation for this ob- servation is given by Proposition 4 in Czarnecki et al. (2017). which provides a necessary condition for function fz) I"" + R" to produce valid gradient vector field. It states that if fa( produces a valid gradient vector field of some potential function ((x), then the Jacobian matrix afe (x) R + R x R must be symmetric. Oz
This can be readily seen from the observation that if de is the gradient vector field to be approxi- Bn
mated by (z) then
0z (z) is the Hessian matrix that is known be symmetric.
A more in-depth discussion is provided in Chaudhari et al. (2024), where Lemma gives a nec- essary and sufficient condition: differentiable function R" + Rn has scalar-valued anti-derivative if and only if its Jacobian is symmetric everywhere, meaning that for all E RR,
ef (z) This suggests good inductive bias should be to restrict fa(x) the class of functions with symmetric Jacobian. This condition is directly satisfied via the Sobolev inductive bias as well as gradient networks proposed in Chaudhari et al. (2024). As pointed out in Czarnecki et al. (2017), this is very unlikely condition to be met by a regular feedforward neural network that would predict the gradient directly as its output.
0f Ba (z)
B.2 DISTRIBUTION PARAMETRIZATION
The Sobolev inductive bias introduced in Section and B.I requires generative model where both the output and its input-gradient are treated as random variables. While the reparameterization trick (Kingmal 2013) could be used with a conditional Gaussian distribution, determining how to distribute the gradient of the samples with respect to the conditioners is less straightforward. To address this, we employ method that relies solely on sampled data and does not assume tractable density estimation of Sobolev random variables.
Moreover, we found that the most common distribution parametrizations for distributional reinforce- ment learning. namely discrete categorical (Bellemare et al. 2017: |Barth-Maron et al.) 2018) and quantile-based (Dabney et al.) 2018h ab, do not scale tractably to higher dimensions, specifically in the way they instantiate Equation g These considerations further motivated us adopt sample- based approach for the distributional Sobolev critic, similar in spirit to Singh et al. (2022); Freirich] generative model that deter- ministically maps noise to samples (Li et al., [2033] Goodfellow et al.) 2020). This parametrization lets model distribution over both outputs and their mpuu-ggradie via the Sobolev inductive bias: differentiating each generated sample with respect to the conditioner yields a reparameterized, sample-wise gradient, which aligns with the intuition from Appendix A
et al. 22019). As result, distributional critic is structured
19

---

Published as a conference paper at ICLR 2026
C SOBOLEV BELLMAN OPERATOR
We first specify our smoothness and boundedness assumptions on the return, transition, reward, and For completeness, we summarize mild conditions under which the neural policy and critic used in our method are Lipschitz in Section [C2] We then derive how differentiating the distributional Bellman equation yields the action-gradient update in Section C.3
policy functions in Section C.1
Next, we show how to bundle the return and its gradient into single affine update rule in Section
C.4
We follow and give both the pathwise and integral forms of the full operator in Section C.5
Finally, we explain how to include state- gradients for full Sobolev backup in Section C.6
C.I PREAMBLE AND ASSUMPTIONS
Unlike traditional Bellman operators acting on scalar- or fixed-dimensional vector-valued functions, our operator acts the space of continuously differentiable, bounded functions with bounded first derivatives over the compact domain x A. Concretely, we assume the return-distrinn diun--iittibution admits a reparameterization
Z (s, az Ez) E CB x A),
where EZ ~ p(Ez) is exogenous noise, so that each sample ((s, a; E2) and its gradients V (s a; EZ) are bounded and continuous on S x A.
Se ~ p(Eg), ~ ple,), and εx ~ p(Ex) be independent noise variables on 2, and
f x Ax A r: r x 2 R,
(reparameterized transition)
(reparameterized reward)
8 x 2 A,
(reparameterized policy)
be C' maps such that, for each draw Ey, Ex;
5' = [(s,a88)) r r= -= a' d'=m(s);8*)) Z Z()
Importantly, we assume that the gradients of all these C mappings (Z, f, and т) are tractable to compute in practice, for instance using automatic differentiation.
We further assume these maps have uniformly bounded Jacobians: there exist constants L. Lif.a Ly, Lmn L2 < such that
of (s. s L..0
81 (s. (s,a; Ls
sup (8,0,8))
sup (e,0,ey)
de
sup (+,0,e,)
8a
Ör (s,
s Tmr,
sup (s.Ex)
8s
0m (s,Ez)
Liz,
V (e, (s,a.z))
sup
S Lz.
C.2 LIPSCHITZNESS OF NEURAL POLICIES AND CRITICS
For completeness, we summarize the mild conditions under which the neural policy and critic used in our method are Lipschitz. Similar assumptions are standard in analyses of gradient-aware RI. methods Oro & Jaskowski] 2020).
Let Fo Rh + do denote either the deterministic policy To(s) the critic Qo a) (determin- istic or distributional). We assume a standard MLP
ho(x)) = 2r, he+1(x) (e(WWheex)) be), 8= 0,.,.., L-1,
with final output
< (WWLh (x) + bz), (policy), w hz(2) +c, (critic).
Fa(x) =
We require two standard conditions: (i) Each activation og and the policy's final activation i is I Lipschitz, with constants T.g and Lag. This includes common choices such ReL.U (Nair & Hinton
20

---

Published as a conference paper at ICLR 2026
2010). tanh, sigmoid, and SiLU (Elfwing al. 2018). (ii) Each weight matrix has finite operator norm [Well2 Σ Mr < oc, which rules out parameter divergence. In practice, weight explosion would lead to immediately visible instabilities such exploding Q-values, divergent TD-errors, and erratic actions, none of which appear in our experiments.
Under these conditions, each layer oeWWe bε) is Lipschitz with constant Lo, Mr. Since compositions of Lipschitz functions produce Lipschitz maps whose constants multiply, it follows that the full network satisfies
[z1) (o(2))
KF 22 x2ll.
with
L-1 Li
(L,
(policy),
KF
||l||2 II (Lo, ML), (critic).
the imput domain is bounded, |||||| S R. and the network parameters remain bounded, then the hidden activations h₂(x) are uniformly bounded on {x ||>||| s R). Consequently, the parameter Jacobian is uniformly bounded, so there exists Ly,0 < co such that
Darammter-Lipschitznesss
sup [|VoFo(>)]| S LF |
In particular, for the policy assume
sup (o)o(()) & L
C.3 DERIVATION OF THE ACTION-GRADIEN TERM
Starting from the distributional Bellman equation for each noise draw (r,
7( = r(s, a;)+) 7 7'',a''ss's a'z)),
where s' =f(s,a;E)), d' = a(s and E'z p(Ez) independently. Differentiate w.r.t.
à r(,,a;E,) + da 0 Z(s, Z(s, da da
d ZZ(s el))
By the chain-rule,
Z(s', "ls',a';e'' d':)'z)) = (fa(s,a;E)) T θ da
d
Os 2(s ('''6)) +
88/00
(f((s, (T=( Z(s', d':6')).
da
Ou' 0a'/0a /0a
where we have introduced the shorthand
dm fa(s, a; sy) da (s,0;e)), T=((';0))= (s':
0f
0s
Plugging back (and evaluating s' f(s, f(s,a;e)),d x(s(s=))) yields
2 i Z(s,
=
da r(s +
da
T 0 las 2(5 2s'''d''''2 2 (T.(s)=e))) & Da
7 (Ja(s,u
Z ((5),')e'2))
which becomes the action-gradient componeen in subsequent block-affine formulation.
21

---

Published as a conference paper at ICLR 2026
C.4 DISTRIBUTIONAL SOBOLEV BELLMAN UPDATE
In Section 3.2 we introduced a Bellman-style update that propagates both the random return and its action-gradient. We rephrase it block-affine form.
Define the stacked random-vector, which by definition coincides with the joint Sobolev return intro- duced in the main text:
(8,, 0; EZ) i da
2 (8,0,77)
(1,))))) R
(ss a; Ez)
where EZ p(E2), and let En be independent noise variables for transition, reward, and policy. For each draw (E En: define
= f(s, Ef); r= r(s, a; Er), d'= T('; Ex))
Then the sasple-wise update is
r(s,a; 6r))
01xm
(())); =
0
+
:(s'd':=2))
aa r(s,
0mx1 ( (s, (0r (s'; Ex)) T
b(s(a18))
A(s,a18/18x)
01xn
((s',''i'')
( (s, a5 a)
Os
N(einsy))
(18)
where for the next-step noise we write (e,, viate
independently. We abbre-
r(s,a;E7) ((8,0,+))= (ar(s,
above.
Next, introduce the state-gradient operator
0
07
D.
c
Os
Os
where ef linear map
(1,0...., selects the return component. For each (s, ayE)), define the combined
a(,,a; + ((s e5)) Da
Then the distributional Sobolev Bellman update can be written as the single (pseudo) affine trans- form
H (s.aa;52) = b(s a; + [(s,a:e),e>) H dl:c
(19)
à Remark (i). this form, the next-state gradient 36 Z(s', s'z) is computed within the update and is accounted for by the linear operator L(s, Ex) via D.. Thus computing the state-gradiens of the distributional return is an integral part of the Bellman operator itself.
Remark (ii). Although we write the update in single "block-affine" line, it is not self-contained in the 2-vector (Z, Z). Indeed, the right-hand side still involves the next-step state-gradient term N Z, which is not part of H. Thus refer this update the incomplete Sobolev Bellman backup. We therefore call the update affine only in a loose, not literal, sense; the genuine finite- dimensional affine form appears once lift to the full Sobolev vector (Z, i, 0,Z) in Section C.6
Measurability. By the same argument used in Appendix A the map (s, 0, E) + θ, (s, a; E) is jointly Bonsl-measurable.
22

---

Published as a conference paper at ICLR 2026
C.5 DISTRIBUTIONAL SOBOLEV BELLMAN OPERATOR
Let S x A P(R+=m) be the law of (s, = (ss a;z2), 0, 0x (,8,;;;2)) 82)], and assume three independent noise variables Ef ~ Pf. Er pr. Ex Pr generate
= [(..0,8))) r= -=)(s,);))) a' = d'= (s)
We define the (incomplete) distributional Sobolev Bellman operator T in two equivalent ways: Pathwise (law-equality) form:
(T n) (s, = Law + ((
(20)
where X' ~ n1s', a') after sampling Er, and setting s' = [(, m(sss)))
Explicit-noise integral (pushforward) form:
( / [ (b(s,a;e,) C(s, xx a
(s, Jr, x p. x (d(n)) x pd(ee))
(21)
C.6 COMPLETE DISTRIBUTIONAL SOBOLEV BELLMAN OPERATOR
In Section 3.2 and above we developed the action-gradient Sobolev Bellman backup, in which our bootstrapped object the pair
7 (s.a) = (Z( 0, E RR+m
Because the Bellman operator needs the next-step state-gradient 0,L(s',a') which is not included in S this update is incomplete. We now introduce the complete operator, which augments the bootstrapped object with the state-gras suate-gradien. Equivalently, we have to own the differentiable function Z itself at evaluation time and thus rely Sobolev inductive bias introduced in Appendix [B.l
Lifting full Sobolev return. Nothing prevents us from treating the sate-gradient as third component of our bootstrapped vector. Defie
Z(s,a)
2 (s,
(g7(s,a)
IR 1+
0,7(s,a)
Then, just as before, we obtain the affine update
S5.. (s, bran (8, a) + Lran(s, a) [ (8'a'))
(22)
where (s,n(s,a) a) E gl+m+n collects (r, Oar, ².r) and Cran(s, a) is bounded linear operator on pl++++n Here Liull is simply matrix, and obtain biall and L by the same derivations in the action-gradient case such that
r(s,açer)
biuu(s, a5E))
0.r(s,ate,)
0.r(s,a;Er)
and
lix1
01xm
0xxn
(:) =y
0mx1
So (5,) (s'sem)
Omxa
Onx1
ff (s, (s(:Em)
ff (s,
With 5' = =f(s, (0,,;ee) and a'= x we denote the complete Sobolev Bellman operator by
(T..., (s, (s,, = Law bran (s, a; ex) + Lran (s, a, [x]]
X' n(s',a)).
(23)
23

---

Published as a conference paper at ICLR 2026
Consequences for implementation. In this complete form, the second and third blocks of S need not be tied to the derivatives the scalar head. One may simply design a network with three output heads
(8, + (z a), (s, Gg(s,a))
and train it against the Sobolev-aftine targets in Equation 12 Of course, in practice one can still exploit the Sobolev inductive bias, thus ensuring G 0,2 and G 0,2.
D RESULTS WITH WASSERSTEIN
In this section, we complete the proof of Proposition 3.1 by deriving the upper bound on the policy gradient error in Section D.2] We then establish contraction properties of the Sobolev Bellman operators in the
metric: Theorem 3 treats the complete operator T while Theorem 4 gives the analogous statement for the incomplete operator that backs up only (Z, 0, Finally, Corollary 4.1 applies Banach's 's fixed-point theorem to conclude that the Sobolev iterates 1/ Th Th converge geometrically to the unique fixed point 17" identified in Lemma
Srremuum-p-Waserrae
D.1 PROOFS CONTRACTION AND FIXED POINT
(209)). Let (X, d) be a metric space and let (X) be the set of all probability measures on X with finite pth moment. For any a, 3.3 Py(X), the g-Wasserstein distance between a and 3 is
Definition 1 (p- (p-Wasserstein distance Villani et al.
1/p
inf [ d(a,v)) dT( (E
Wy B)
xxx
where IT 3) is the set of all couplings of a and B, Lee Borel probability measures X x X whose marginals are a and 3, respectively.
Definition 2 (Supremum- p- Wasserstein distance). Let (X, d) be a metric space, let S be the set of states and A the set of actions, and let M2: A + +P((X) assign to each action pair (s, a) probability measure on X with finite pth moment. The hupremum---aaseee between μi and Pz is
distance
W (2)
sup (s.a)e8xA
(p(ni(s, u2(s,a)),
where Wp the p- -asserssstnin distance on (X, dl) as in Definition
Lemma (Push-forward law identity). Let Z be a random variable with distribution μs, and let f be any measurable function. Then
f4u = Law (fw()(Z)).
Proof. For any Borel set A,
Pr( E A) P)) = = f*M(AA)
Since this holds for all A, conclude f# Law(f(2)
Lemma 2 ((Pseudo- Jaffnne pushforward form of the Sobolev Bellman operator). Fix (s,a), a). Draw independent exogenous noise
Ey p(Er), Ef p(Ep)))) en p(ea),
and set
r= s' f(= d' d'=e(s)eex), x' x'mrss(ss'a) Define the random (pseudo- Jaffine map
( (
Va Er, Er,E)) = b(s, + ((s) where b as in Equation
L C
dar)
19 Then
(T (5) Law (D..a (X)).
Law
In other words, at a fixed (s, a). conditional on the exogenous noise, the Sobolev Bellman update is affine pushforward of the joint return-gradde -raud-gradien...
24

---

Published as a conference paper at ICLR 2026
Proof.
Fix Borel set A R 1 Writing out the randomness explicitly in terms of Er, ExxEx
have
Pr( E A) = Ber.syien (Pr( (X')) E A
6r, Er,
Conditioning on (Er, ex) turns 2s,a into the fixed (pseudo-Jaffine map a b + C[x) so
Pr(0.,d (X') A Ex, Pr(b + f|X* A). Since X' (f d'), an application of Lemma |
((b+ [bb+ " (s', a') Law(0 (b + [(X')
gives
Pr(b -[X E A) = [(b + L )) # n (s' a (A)]
Putting these together,
Pr (X') E A) = Esnre Fer (b+L) 11 Sm (s', d (A)]
Finally, by the explicit-noise integral form Equation 21 this expectation exactly the definition of (TS s (s, (A). Hence
Pr(Os.(XX') E A) (I 17 (s,
and since this holds for every Borel A, we conclude Law aa (X') = (T ) a).
Lemma 3 (Affine push-forward contraction). Let (X, d) be a normed vector space equipped with the metric d induced by its norm, and let
F: X + X,
F(x) = b + L(z).
where b X is fixed and L: X X is bounded linear operator. Define its Lipschitz constant
d(L(ar), Lwwl)
| sup xFy d(z, v)
30.
Then for any two probability measures 3 on X with finite pth moment,
(F(F+a, F사 S |[lla Wy 3)
In particular, when L 7 this recovers Lemma 3 of Zhang et al. (202)
Proof. Fix any & > 0, and choose coupling
क II( 3)
that is e-optimal, meaning that its transport cost is within of the infimum:
( d(z, y) (x
1/p
Wa B) +
Push this coupling forward under /' F to obtain
n' = (F x F)+T [F F사
Then
W/ (F P4 < 1 d(u,v)' d=' (u,
(by definition of Wp)
xxx
-f xxxx ((F(z),F())) dπ(
(I' (F #)
-f XxxX d(L[x), (z(
(F(x) F(y) L[x y)
|| I
d(z, 1) dz(
(d(L(x), I(v)) S | d(x, v))
Xxx
(V (3 =)²
(m e-opimall).
Taking the pth root and letting E + yields
W, (F F+ 3) < ||a|aa W,(a,))
25

---

Published as a conference paper at ICLR 2026
Lemma 4 (Mixture non-expansion (conditional form, Zhang al. (2021). Lemma 4)). Let C' be a random variable on (S2 F, p). and let Z1, Zg be I2 -valued random variables. Let p 2 and suppose there exists ò 2 such that for every E 52
WL(aa (Z1 C CC Law (z CC= c)) S 8.
Then the marginal distributions satisfy
W (Law (21), Law (22)) 2 d.
Equivalensly,
sup W, (LawWWZZ -==C= Law (22 C - c)) S cEft
(L(aa (Z1), Law (22)) S 8
In other words, averaging over the conditioning index cannot exceed the supremum of the condi- tional Wasserstein distances.
Proof. This proof follows Zhang et al. (221)) Fix any 0. By assumption, for every Ω
W, (Law( Zi C- c), Law(2) Z2 C =c)) S 6,
so there exists a coupling
C -c), Law (22 C = Ccc) To II(Law |
such that
J 2xdx d(z, y) d 8) - (6 E)P.
By the law of total probability, the marginals are
Law (aw(Z))))= = Law( [1 C = c)(-) pldc),
Law (aw(Z))())= Ja
Law (Z 1== ()) p(dc).
Define global coupling म on Idd R by
*(U)= TC(
U UC CCR x Rd.
nn
Then for any measurable C Rdd.
/ (c RRdd p(dc)
# {AA
Law (Z C C= (A) = (A))
ja
and similarly ग (Rdd A) Law (A), so 7 E [[[aaw(()) Law((z))) meaning a valid joint law for pair whose marginals are Law(Z, (2) and Law (Zz). Hence स is admissible coupling in the definition of W (Law(Z))) Law( (2)).
{ Rd
W,(LLW(Z,), Law (Zz))²
d(z, v) d= y)
(by definition of W,)
[ [f d(s, u) d (x. v) pldc)
(by definition of ボ)
<
(6 +e) p(dc)
(we is s-optimal)
Jn
= EPP.
Therefore, and since E > 0 was arbitrary, letting - 0 yields
WLaw (Zz), Law (Zz)) < (6 +
(y(wa((Z,) Law (Zz)) < 6.
26

---

Published as a conference paper at ICLR 2026
Lemma 5 (Mixture p-convexity for Wp). Let (X, d) be metric space, E [1, oo), and let (S2, F, ,)) be a probability space. Let (He)cen, (Volcen C (X) be such that the mixtures Jq Pe p(dc) and Ja Ve p(de) are well defined. Then
(1, Wolleseed olde)
1/p
wo( J/ ne p(de),
/ vopldes)
Proof.
Step I: s-optimal couplings for each c.
Fix > For each E Ω2 pick an e-optimal coupling Th E I1 Ve) such that
J XxX
d(z, u) oe (dx, dy) Wp Vo) +e.
Step Measurable selection and mixed coupling. Assume the couplings x. can be chosen measurable function of so that the following mixed coupling well defined
#(U) S [x(U))
U. Xxxx X Borel.
For any measurable A A
#(Axx)= J TT((A x) p(dc) / Mc(A)) (L/ Mc pa (A),
and similarly
# x / (x x A) p(dc) = J, vo(A) vo(A) pldc) ve p(de)) (A). (
Hence 7 has the mixed marginals Sa Me p(dc) and Sa p(dc), meaning that I Mc I ve pellec)).
т E
Step 3: Bound the transport cost of the mixed coupling.
Since (c, x, y) + d(x y)" is nonnegative and measurable, Tonelli's theorem allows us to exchange the order of integration in (c, x, v);
fxX d(x, u) #5 (dr, dy) = Jx xxx d(a,y)? (L (frilde dy) p(de))
J (.Ver d(r,v)) (d dy)) pldc)
(WolHes ve) + +e) plde)
[ Wolues u. pldc) + e.
Step Take the infimum over couplings and pass the limit. By definition of W.,,
whf |e dp, { [e dp) < { d(r,u)P &5 (dx, dy) s J/ WplMes uc) p(dc) + e.
Taking pth roots and letting - 0 yields
1/8
Wº( wa(/. Me p(de), ( Ve p(dd)) s (/Woleetes p(de)
27

---

Published as a conference paper at ICLR 2026
Lemma 6 (Spectral norm of block-triangular matrix). Let A Rmxm B E Ruxm and C Rax" Then
0
A
< max 1, 2, [[ |
B
C
1
Proof. Split
1
0
0
0
0
M =
0
A
0
=
0
A
+
0
0
0
0
B
C
0
B
0
E
By the triangle inequality for the spectral norm,
M & ||]|2 + F1
Since D is block-diagonal, |||| = (1, [1, |1-112, |(C2))) and E has only the single nonzero block
B, so | B1|2- Substitution gives the claimed bound.
Theorem (Supremum-y-Wassee Wasserstein contraction of the complete Sobolev Bellman operator). Let
T (8x A P(R
A
be the complete Sobolev Bellman operator bootstrapping the vector (Z, (2.0,2,0,2),
The distributional Sobolev Bellman operators are defined by the reparameterized policy x (s'; (s'(:x)) 1x)) the transition [s,,;E E,) and reward r(s, &; By assumptions from Section C. we write
||A|| sup sup [[ss, S (max(1, Lg,a Ls + L L = y Kfulls 8.0. 5y.t
then.for any two Sobolev cuuum-distribationfuuncitii
122
< || Wy(m v2) S 7 Knan Wy(m, 12). W W mn, 112)))))
In particular, Th a | whenever ||A||a < 1, and sufficient condition for this 7 Kiall < 1.
Proof. We show that the Sobolev Bellman map T is || conttraionn in the supremum p Wasserstein metric.
By definition we have,
W,(TS-m T Tg sup W (s,a), T m))))))
Let's fix arbitrary pair a). Then draw the exogenous noises for both updates
E, ~ p(rr), = En ~ p(=x),
set
r= == r(,,0e))), s {(, d' = m({);e)),
and sample
X ~ mi
Sa. (s',a'),
x² ~ S.a (s'a'')
Define the random affine map
(I; Er, En,8)) = b(x,05e)) + ex)
r
where b =
0ar
L =
0
fT π T
0
0,r
f T
fT
28

---

Published as a conference paper at ICLR 2026
in Equation 22
By Lemma we have
T (s = Law Law((ea(X;)) i= 1=
so
WA (8, a), m2 a)) = W "LawLa(B (o Law((D,, (X2)))
a'), for every (s',a'), Since Law(X() 11i Sr. (s' (s',
122 San (d',d)) S WA(m, 12). Wx(Law(Xi) Law(X))) - Wy(th Sa.e (s',
Condition on the full noise triple (Er, Es,))) so that Lsoa is a deterministic affine map. By Lemma W Wy(Law( ( aa.a (X1) er, 6,, En); Law(D..(X)) ler (r,)),
(s) W Whnh "n"""ss a'), 1/2 , Wa 72).
Sa. (s',
Taking the supremum over e,,e, and then applying Lemma yields
W WyLLaw(Qs,a Law( s sup en)) Wynn, m))
Ey.5x
Using T (:(s,0) = Law(ls,))) (X:)), we conclude
Wp WT (s, a), T na(s, a)) = WW Law(Is,
I sup [1 ((), Wy(n, 72).
Egien
Finally, taking the supremum over (s, yields
W, (T S M11 mp)) = sup W, (T (8, T 12(s0))
a.a
< sup sup I[[(s (s,,,,,,,, Wy(n, 72)
Efie -Illla Waln m)
T
is contraction with coefficient [[]a.
By Section C.I
we have
(s. a,y))ll S Ls (s, el s Lg.:
[Ts(d\i)))ll s Lx.
Hence for each (s,0,8y,E*) and applying Lemma
6 we have
0
[[s (((,,,
0
fa
T
0
0
S
T
S
> max 1, llft o. 112. I1 l12)
+
1
1l
diagonal blocks
off-diagonal block B
S 7 max 1, [|fallalTTolla [[falla] + [[fallal[=alla)
y maxx(1, Lg,Lxx, Lj.e} L1,oLa) y Kiull-
Since the last inequality holds for every choice of taking the supremum gives
[| sup [[ Y Kfull-
s.a.ef.t
29

---

Published as a conference paper at ICLR 2026
Lemma (Fixed-pont fixed-point law of the complete Sobolev Bellman operator). Define the infinitehorizon return and its full action- and state- satte-gaddeens under policy π by
0 0g > Tt
Z(s,a)
Σ
Tt
(s, = 0
7 Tb.
G"
Σ 1=0
and let
- Law a), G" G*(s,))).
Then 1f" is fixed point of the complete Sobolev Bellman operator:
T r7" 17"
Proof. Recalling the one-step affine update
@ a E,, Ex) binn(s, + a;
the Bellman recursion and its derivatives combine
(7(s, a), (" (s, $s.a a') G" ('))'
By Lemma 2 we have
Law Law(Z(s,o), (s, ((s.) = Law P a (S)) -T 17" (s,a).
Since this holds for every (8, a), we conclude T 11" = =
Theorem 4 (Supremum-y-Waseestid contraction of the incomplete Sobolev Bellman operator).
Let
T (S x A P(R!+ - A (
be the action gradient Sobolev Bellman operator updating only H(s,o) (Z,s,a),0,Z(s, Let p 2 1. Fix any two return-grddiet laws M,, 72 S A P(g)+mm) Assume
[alll S Lfgas
|| S Lu-
(Reparameterized pathwise gradients and coupled lifting). For each (s', assume there exist random variables
(H' GR X (H2 G2) E R 1+ RR.
defined common probability space, such that H: 0, O,Zi(s', d') is the pathwise (reparameterization) derivative as in Appendix|A and is induced by our distribution parametrization (Appendix B.2). Moreover, assume that there exists optimal pling C* E [[p (s' aa'), 72(s', (.'.a. a')) such that (H(, H2)) ~ S" and the joint law of (H) G H3, Gy is compatible with 5", meaning that its (H', -marginal equals Ç* and it couples the correspond- ing pathwise derivatives Gi, G2
0,4) ()(s'aa)) and G¹
Assume there a constant LD, 2 0 such that for every (8,,a) aa)) for the above coupling,
E| E 111G) 1/8 s LD. E [EIH H3 1/p
(24)
Set
= (1, Ls Ior] + L.f.a LDD,-
Then
ST m. T 92) S Ketf Wp 12).
If Keff 1, Sa is strict contraction.
30

---

Published as a conference paper at ICLR 2026
Proof. All vector norms I|lI| are Euclidean norms, and operator norms are the corresponding induced norms. All Wp distances are taken with the metric d(z, y) = || y|l-
Fix (s, a) and draw sample (Er, Ey,x),, Ex). Set (, and write
= [(x (ze)), a' = ;
and define the (pseudo-)affiin Jaffine update acting on the extended pair (H', G') by OHH,,") = b(s, (,,) + A(s, Ex H' + N(s,
For this fixed noise draw let The denote the corresponding one-step update, so that (s, a) Law( G²) c), where the remaining randomness is only that of the next- step draw H' n d') and its attached pathwise gradient G".
By the uniform Jacobian bounds of Section C.I
sup [s, Ey))l < Lf s,a,f
sup || I La, s',5x
so
(() max
(s s maxx1, Lggaber
and similarly
(s = (I 0z o3 s Iif.a-
Let (Hi, Gú, Hzy G%) be coupled as in the theorem assumptions, so that (Hi, Hb) ~ C is an optimal coupling of 7/2 (5', a') and (2 a'). Write H = X e 1 Then
[[ G)) (X2, G1
|A(Xi Xx) + NGG G2)
(definition of Φ)
< |A(Xi X3)II + |[N G/)l S ||||| - X121 + |NN -
(triangle inequality) (operator-norm bound).
Taking pth powers, expectations, and pth roots, and using Minkowski's inequality yields
EI G1) - (X2; 6)) 1/P
I (E(IAIIII - +51++INNIIG -GI) 1/P
(raise to p. take E, take pth root)
< |IAI -XII) 1/P
+ [IN][ (E|G -BIG)--111)
1/r
(Minkowski inequality).
Invoking the derivative- coupling bound Equation 24 gives
-BIIXiXx111 2/P
[B| G 4(22, G(2 < ([[A]I
Since (X1, Ç is optimal,
[BI x211") 1/2 WV 12),,'(
Define (U.V) == (x(,,, G1), (X2, G2)).. Then Law Law(U, (U,V) is an admissible coupling of the corresponding one-step updated laws noise draw c, hence
W. (8, a), (s)) s (BLJU vII))"
1/P
Combining the last inequalities and substituting the bounds on ||A|| and ||]] gives, for this fixed noise draw c,
W W a), T Te a)) s (max Ly.bax + LIaLD.)) Wo(miss' (a),s,a),,))) m'), Kett Wy(m(s'' a'), ()(s)sa'')))
31

---

Published as a conference paper at ICLR 2026
By Lemma 5 mixture p-convexity), and using that Wp Wp(m(s', .'), me(s' = (
" ( (=))a( Ee
W ( (T))
(7 Kieft 1 Ee W7(mm (5', a'), na(s',
(7 Kett ? W (m (m, 72).
Taking pth roots and then the supremum over (s,a) a) yields the claim.
Lemma 8 (Fixed-point law of the incomplete Sobolev Bellman operator), Define the infinite- horizon return and its full action -raionnn under policy π by
Z(s,
G(s,a) 0s
7
Tt.
1=0
1=0
and led (s,( a(), Then 1f is a fixed point the Sobolev Bellman operator:
T 17" = 17"
Proof. Recalling the one-step affine update
& (z; Er:Ey, = b(s, + =(,
the Bellman recursion and its derivative combine to
(Z(s, a), G() $..a (Z(s' G(ss,a')).
By Lemma 2 we have
iiss',o aa), Law (Z(s, G(s, a)) Law (Z(
Giss' d'))) = (T (s.
Since this holds for every (s, a), we conclude 1 =n=
Corollary 4.1 (Convergence of Sobolev evaluation iterates). Under the conditions of Theorem let denote the contraction constant with к = Reff for the incomplete operator or x = Keull for the complete and suppose 7 K 1. For any initial Sobolev return-distribution function V)o. define the sequence
Mh+1 = T Tm.
where T may be either the incomplete (S,) or complete (S..a) Sobolev operator. Then by Banach's fixed-poin theorem the iterates converge to the unique fixed point 17" (ef. Lemmas 7 8
W, (T,, 7²" < (7 k)" Wp (10o0o n7*)
12+00
particular, the + 17th the supremum Wasserstein metric.
D.2 PROOF OF PROPOSITION 3.
Lemma 9 (Mean-difference bound via Wi). Ler X, be Ré--vlued random variables with distri- butions = Law(X) and v Law(YY, and assume E X|| < oc, E|||| < co. Then
[[ - [||] S Wi(,,v),
Here, |- denotes the Euclidean IRd, and W1 is taken with respect to the ground metric d(z, y) = || vll-
Proof. Let
mx = E[X]. my Ely]
mxx * my, define the unit vector
mx my
a
mx my
32

---

Published as a conference paper at ICLR 2026
Then the scalar function f(æ) = u satisfies
- -f(z)- = | (z - f)1 - || || - -yyl = || - y|l,
so tpp S 1. By Kantorovich-Rubinseri duality (Villani et al. 2009).
Wi(k, sup
11 [[][ E\g( 2 [[[(()] [B][()]]]
But [[f( [[f(Y) = u (mx my) = |mx - myll. hence ||mx - my I W1 v).
mx my, the inequality is trivial.
Proposition 2. Let a be an Iippo-iip continuous policy, and let Law V 7. (x, la=x(e)) and Law L..[V.2(s... a) la=a(s)) denote the true and estimated distributions of the action-gradien at (s). respectively. Then the error between the true and estimated policy gradients bounded by
VaJ(0) - [[] S
Lx.00 Bondy (L a) la=+())l, Law a as) la=(
I I-7
This proposition is a distributional extension of Proposition 3.1 from D'Oro & Jaskowski 2020
Proof.
Step 1: True and estimated policy gradients.
The true policy gradient is:
VoJ(0) 1
E (F VoT(s)]
1- Bandy
and the estimated policy gradient is:
voJ(0) = sn
1
E
V. V. a) (==((e) VoT(s)
Step 2: Policy gradient error. The norm of the difference between the true and estimated policy gradients is:
VoJ(O) voJ(o)| (v)00)
= | Ba-az (E (s,a) - - (V -
x ( (s)
Step 3: Applying the Triangle Inequality and Lipschitz Continuity
Using the triangle inequality and the Lipschitz continuity of the policy ([Vo I Lx,e), we have:
VoJ(O) -oJ0))-aa(001 S -
1-ク Exndz
E Vo (8,, a) [a= E (v.)(s, )
x [
< L B. Bandy
|| V.2% (s, 05, [=mto)) E v.2(s, a) [ammme
Step 4: Bounding the mean difference by W1.
Let
= Vx (s,
Y= V=
By Lemma 9 we have
[[ E[Y]|| < Wi(Law(X), Law(V))
Wi(Lave (o2" (s, a) ) Law Law(Va V. (V a=7()))
33

---

Published as a conference paper at ICLR 2026
Step 5: Conclusion
Combining the results from the previous steps, we established that the L2 norm of the difference between the true and estimated policy gradients can be bounded as follows:
VoJ(0) - voJ(O)] -
La. 1-7
Bondy W W (Law| 2= la=x(s)| Law(V ((s a) le=
34

---

Published as a conference paper at ICLR 2026
E PRACTICAL DIFFICULTIES WITH WASSERSTEIN FOR TRAINING
Adversarial W1 a loose proxy for true Wasserstein. Some works have directly applied WGAN training to distributional RL, for example Freirich al. (2019), who cast return distributions as largets for adversarial matching, In practice, WGAN training replaces the exact Kantorovich- Rubinstein dual (Arjovsky al., 2017 Gulrajani et al. 2017) with a parametric discriminator plus approximate Lipschitz control (weight clipping or gradient penalties). This induces three sources of deviation from the true distance: umite-capacity approximation error, imperfect Lipschitz en- forcement, and optimization error. Systematic analyses show that the resulting WGAN losses can correlate poorly with the actual Wasserstein metric and need not be meaningful approximations of it, undermining proofs that presume access to the exact W1, See both the empirical and theoretical analysis of Mallasto et al. (2099) and the critique by Stanczuk et al. 2021
Computational cost of exact OT in multiple dimensions. Even ignoring estimator issues, com- puting multivariate W, exactly on mini-batches costly: building the pairwise cost matrix requires (m memory, and solving the discrete OT problem takes ((m² log m) time (Genevay et al, This makes per-update calls prohibitive in distributional RL..
2019
Sample complexity. Beyond runtime, OT also suffers from poor statistical efficiency: the em- pirical Wasserstein distance converges to its population value at rate O(n-1/d) in dimension d, compared to O/n_ for MMD (Genevay et al. 2019). This slow convergence further limits its practicality in high-dimensional RL.
35

---

Published as a conference paper at ICLR 2026
F BACKGROUND ON MMD
Definition 3 (Maximum Mean Discrepancy as an IPM). Let k: X x X - R be a symmetric, positive-semi-defiiee reproducing kernel with RKHS H and feature map
Φ: X H, Φ(x) k(x,,))
For probability measures P. on X, define their mean embeddings
uQ | ((x) dQ(x).
Pp ((x) dP(x),
Then the Maximum Mean Discrepancy is
MMDx(P,Q) = ||p MQlwx
whose square admits the kernel expansion
(P,Q) 1|pp Me l
// h(2, dP + aa)
2 2/ʃ (P())
/I MZ,U))
(v)
Moreover, MMD₁ coincides with the integral probability metric (IPM) over the unit ball of H. namely
(M.)(P.Q) = sup [Fa-r[f(>)] Bg-olf(u)] = ||r pQllx.
JEH 11 S1
as shown Gretton et al. 2012
Remark (Euclidean densities). Working in Rd, if P and Q admit densities p(x) and q(x) with
respect to Lebesgue measure dr, then dl(x)= p(x) dr, dQ(x) g(z) dx, d(p and each of the above integrals becomes an ordinary Lebesgue integral in x. Definition 4 (Conditionally positive definite (CPD) kemel). Ler X be measurable space and let k X X + R be symmetric. We say that l is conditionally positive definite (CPD) if
= (p(x)) g(x))) dz,
JJ xxx
M'xxx') dp(x)) dp(x') 2 for all finite signed measures on X with μ(X) = 0.
If the inequality is strict for every nonzero such p. then k is conditionally strictly positive definite (CSPD).
Proposition 3 (Equivalence of 7 and RKHS-MMD for CPD kernels). Let k X x X - IR. be conditionally positive definite (CPD) and define
ps(x, y) k(x,x)) k(y, y) 2k(x,y))
Fix 20 E X and set the distance- (one-point centered) kernel
1º (xx v) [px (xx 20) Pk (y, 20) PI( v)] k(r, y) k( z0) Rs 20 v) k(20, 20).
Then k2 is positive definite and admits an RKHS Hyo For any P, Q with finite integrals,
(P) = // k( v) d(P ()) 2)(2)
/ʃ h (((x,)) PPP))
||uee (P) - 1p (Q) ||2 Hen
= -MD (P.
36

---

Published as a conference paper at ICLR 2026
Justification. This follows from the distance-induced kernel construction and equivalence results in Sejdinovic et al.) (2013).
Proposition 4 (MMD as a Metric on P(X)). Let k: x X + be a symmetric kernel. We say that MMD, defines metric P(X) iff k is conditionally strictly positive definite (CSPD), ie., for every nonzero finite signed Borel measure v with v(X) = 0.
J| k( v) dv(xx (z) > 0.
Then MMD; satisfies the metric axioms on P(X):
1. Nonnegativity: MMDD&PP.Q) 2 0.
2. Symmetry: MM, Q) = (M,
3. Identity of indiscemibles: MM,x(P,.0) = =0 P=Q.
4. Triangle inequality: for any P. QQR. R E P(X), (MD+(P.Q) MMD R) + MMDx(R,Q) Q).
Justification. This is the standard correspondence between negative-lypp distances, distance- induced kernels, and RKHS MMD metrics as outlined in Sejdinovic al. (2013).
Examples of kernels inducing metric.
Gaussian RBF kernel: for any bandwidth o >
k IRBE' (z. v) exp(-1|z -0l2/(20²))
which is characteristic on Rd and hence induces metric on P(Rd) via MMDRRBF. KRBF.
Multiquadric kernel (Killingberg & Langseth (2023)):
k
MQ (I,Y) - V ++ 122 1 U
h> 0,
which is conditionally strictly positive-definite and thus induces a metric on distributions via MMD, MQ.
F.1 CONTRACTION UNDER MMD
Contraction guarantees under MMD can be established in much the same way as in Appendix This first requires defining the notion of the supremum-MMM which, as with the supremum- Wasserstein distance, is worst-case bound over the entire aate-action space:
D
MMD (n. v) = sup MMD(1(s, a), v(s,a)). (SxA)
As Nguyen al. (2020) first introduced MMD-based distributional reinforcement learning, they provided criteria under which a kernel induces a contraction in this sup-MMD metric with the stan- dard distributional Bellman operator. We first recall the univariate distributional Bellman operator Bellemare et al. 2017):
(iat (s) Laaaaa(sss)) a) (7
~'^ P(. s..).
Sufficient conditions. Let k(x,y) = E es k(( a positive-definite kernel on R. each component k; satisfies:
1. Shift-imvariance ke(x + c, (a for all e R.
2. Scale-sensitivity of order aş: ka(ex, cy) |d|a ke(x.v) y) for all R,
37

---

Published as a conference paper at ICLR 2026
then for any policy n,
MMD no (T 17. T S $0.//2MMD MMD.0
where a, = minger og.
Contraction under the multiquadric kernel. Another work, whose kernel we use, Killingberg & Langseth (2023) proposed the multiquadric kernel discussed above. They showed contraction in MMD" for that specific kernel. More precisely, this holds for any pair of distributions μ4, P(R),
letting (Jrr) # ( denote the pushforward of ja by + r + 722, they show
MMD ( (fro) #v; k)) I y MMD (U, MN))
Taking square-roots on both sides gives the pointwise MMD bound
MMD(UFnn) (Fri)#y; 1)) < v> MMD(U, kN))
Define the supremum-MMD over state-action pairs by MM (1) v) = sup MMD()(s, a), u(,,a)) (x.
Then for any two return distribution mappings n,v, 17, T Tr
MMD (Triist,,
sup (MD(( (#Vl.)) (s,a)
S sup v7 VTMMMDn vss,a)) (4,0)
v Mo (n.
Thus the distributional Bellman operator is - conttraaat in MMD
Contraction in the multivariate setting. To the best of our knowledge, the only MMDDccotraction result in a multivariate setting is from Wiltzer al.) (2024). They show that when each sampled return vector Z E Rd is pushed through the affine map - R(s,a) y2, the same two re- quirements, chift-invariance and homogeneity of the kernel, guarantee 70/2 contraction in the supremum-MMD metric. In other words, by treating each component of the return vector uniformly and applying the identical homogeneity-band -baed argument from the univariate case, one obtains ex- actly the same geometric shrinkage factor.
This, however, falls short of the setting we require for the distributional Sobolev Bellman operator in Appendix C where the pushforward the more general (pseudo)-affine map
2 Ds.a (x)) b(s,a) a(s,a)[=z)
and C(s, a) need not be simple diagonal scaling. Characterizing the precise conditions on both the kernel and the (pseudo- linear operator C(s,a) a) under which this general &s.0 yields contraction in supremum-MMD remains an open problem.
F.2 EMPIRICAL ESTIMATORS OF MMD
In practice, the expectations in Equation 17 cannot be computed exactly and must be approximated from samples. Given two sets of m samples {ax mm
ri P and fu ~ Q. one commonly used
estimator the biased form (Gretton et al. 2012):
MMD = m²
Σ +J=1
k( 2T)) + m²
k(y,
k( Uz).
(25)
m J==1
JJ=1
An alternative is the unbiased estimator proposed in (Grtton et al. 2012):
2
MMD m(m
M(0,21) m(m
k(y. UU)
k(a,, vj).
(26)
m²
The biased estimator given above is commonly preferred over the unbiased one in works involving MMD for distributional RI. (Nguyen et al. 2020 Killingberg & Langseth 2023).
38

---

Published as a conference paper at ICLR 2026
G RESULTS WITH MAX SLICED MMD (MSMMD)
In this section, we study max-sliced MMD contractions for Sobolev Bellman updates. Theorem
gives a contraction result for the complete Sobolev Bellman operator T under the supremum- slixes MMDs discrepancy, under appropriate conditions on the underlying kernel Theo-
rem establishes an analogous contraction statement for the incomplete action oction-gradiens operator
T under distributional lifting assumption.
Lemma 10 (Scale-contraction of MMD under the multiquadric kernel). Let h > O and consider the (negative) multiquadric kernel
ki (2,, y) -V ++ - 2. E IR.
For probability measures J.v on IR with finite second moments, define
MMD (1,2) =
- + -
for X, X' ~ i.i.d. and Y, Y' v i.i.d.
For the scaling map S x + with E [0, 11. we have MMD (H,
MMD K,, ((S.) (S.)
05851.
(27)
Consequently,
M.DD,.( ((S.) #1, #V) S MMD (H.
05s51.
(28)
Proof. If s = O, the measures collapse to Dirac deltas at the origin, yielding MMD zero on both sides, so the bound holds trivially. We assume < & S henceforth.
For any u > 0 we use the identity (see (Schilling et al. 2012 Chapter 15.2, pp. 218-219)):
e (w)
3/2 dt.
(29)
2/T Jo
Multiplying by -1 and substituting a 1 11 y112 yields the representation:
1
100 (e-1(++b"l===11")) 1) 3/2 dt.
kn(
(30)
2
Substituting Equation 30 into the definition of MMD leads an expectation of integrals:
MMD (H,
1
(e ((((++)I--XX 1) 332 dt
2
(e (1 1) 1-/2 dt
2 (e-((++P)X-X1))) 132- dt ]]
To interchange the expectation and the integral fo dt, we invoke Fubini's theorem. For this, it suffices to show
-1+++1 3/2 di < oc,
since the two other terms treated analogously. We split the integration domain fo So + f and bound each part separately:
. For 0 < t < 1: using the inequality ==- - 1| S 2 for 2 2 0, we obtain
le- 1| S tu,
for any 2 0. Thus the integrand is bounded by a constant multiple of 1-1 Since S/ 1/2 dt < oo, the integral over (0, 1] is finite whenever E[u] < co. the present setting, for the first term we have = 1 + 12 X - x''12, so
Elu] = + h E x X'112 < oo,
which follows from the finite second moments of p.
39

---

Published as a conference paper at ICLR 2026
For >> 1: using le- $ 1 for 2 2 0, we have
e le- - 3
so the integrand bounded by constant multiple of 1-3/22 Since f/ 1-3/2 di < oc, the integral over [1, oc) is finite.
Hence
- dt dt= / (...) (...)) dt+ f((..)at oc,
E
(...)
so the absolute integrability condition is satisfied and Fubini's theorem justifies the interchange of expectation and integration:
MMD
Kn, (H,D)
(1 -1 -- dt.
2V
2(e -1+ 1)
A(t)
The constant terms (-1 1 2) cancel out in A(t), only the exponential parts remain.
We now introduce the Gaussian kernel
(xx,01 v) = e --11===11
> 0,
and its associated squared MMD
()(H)) E (X, X') ERC (X,,
With this notation, we can rewrite A(t) as
A(t) e- <Be
Ee Ee- -Y -2 2E -2E--12(x- = :
This yields the integral mixture representation:
1
( 3-12-3/2 Dig v) dt.
MMD (H,
(31)
2/
Now let Me (S+) #4 and (5.)4v. The MMD under the Gaussian kernel scales
(76 We apply Equation 31 to Ma, and make the change of variables u = ts?. Then u/s², di du,, and the measure transforms as:
(7, Ps, v.) = (yys μr v).
3/2 dt (u/s²)-03 3/2 2 3/ -du su 3/2
Substituting this yields:
MMD2, ky (M., (p))
s u1/s? u 3/2 -3/2D2 (uh² He v) du.
(32)
2/T
For 0 < S 1, we observe that 1/5² 2 1, which implies the pointwise bound:
u//² S e for all > 0.
Since the remaining integrand factor 3/2 DE (uh²: v) is always non-negative, this inequality is preserved upon integration. Combining this observation with Equation 32 gives
13/2D2/
MMD ky, (Pa: V/)) I
2/7 e
3/2 (uh² μ. v) du.
(33)
Comparing Equation 33 with Equation 31 yields:
MMD? (Mo, v2) < MMD2 (M.
Taking square concludes the proof.
40

---

Published as a conference paper at ICLR 2026
Lemma 11 (Mixture p- conoexiitt of MMD; in RKHS). Let x R be symmetric Setriiee-emddeffiiie reproducing kernel with RKHS (H, (.,)) and feature map 6(x) k(z, .). Let F. p) be probability space, and let (Ne)een and (Ve)eeR be families of probability measures on X such that the mean embeddings Pęc == Jx & djte and Pve == Jx & due exist in H. Define the mixtures p == So Pe p(dc) and & == Ja p(dc), and assume that up, pe, Sa Mu. p(dc), and Sq Mv, p(dc) are well defined in H. for every p E (0)
MMDx D) - MMMD v.) p(de) 1/p
Proof. By linearity mean embeddings,
Pr = S/ Mnn p(dc),
Pe m==1 Mv, p(dc);
where Pu: = Sx (xx) (x)(x) and Mve - fx 6(x) duc(x) are elements of H. Thus,
Na ps / v(c) p(dc),
v(c) = Mpe Mue E H.
12
Hence
MMD+(R, = 11m moll2 = 11/ o(c) p(de))
$ ( [o(c))ll
(triangle inequality in H)
< (L [(((e))l p(dc)) 1/D
(L' L a probability space).
Finally, [[b(c]] = l1ppe - M MMD (Me, which gives the claim.
Lemma 12 (Mixture p p-convexity for CPD kernels via the distacce-idducd RKHS). Let k X + R be conditionally positive definite (CPD) and let k- be the associated distance-induced (one-point centered) kernel from Proposition so that for all probabilities P, Q with finite integrals,
"ik (P, Q) = MMD (P,Q).
Let (..FF.p) p) be probability space, and let (pe)ees and (Ve)ees be families of probability mea- sures on X with finite embeddings for kº Define the mixtures i := Sq He p(dc) and D = Sa Ve p(de), and assume that p and i also have finite embeddings for ko Then for every E [1, oo),
(nanlee Dc)P p(de) 1/p
7k (p, =
Proof. By Proposition 7/ MMD Applying Lemma to the PSD kernel ke and the families (He), (vo) yields
MMD (7, D) & (/ MMD (Me, v.)p p(de)) 1/P
Replacing MM by 7x via Proposition gives the claim.
Lemma 13 (Max MMx--siced affine push-forward contraction for MMD3). Let k: I x R R be kernel on IR, and let MMD; be its associated maximum mean discrepancy on P(R). Assume that for all μ. P(RR):
T Translation invariance: for every t R.
MMD ((x: f)
MMD v)-
(S) Scale-contraction: there exists a nondecreasing c [0, + [0, oo) such that for every
s e [0. 1).
MMD& + 8x) (xr sT) S c(s) (M,
41

---

Published as a conference paper at ICLR 2026
Define the max sliced lift of M by
(SM v) == sup (Moo((Po) (P))#v), Pa(x) = =0, 0E32-1
Let F(x) = Ax with an arbitrary matrix A Rdxd and E Rd, and denote == ||Allop
s1llli|==1 |ull If E [0, 11, then for all μ1, U ((TR))
MSMMDA (F#Hu FWV) S cL) MSMMD (1,
Proof. Fix 0 E Sd-1 and set wh == ATB.
Case 1: we 0. Then (Po (x) (8, b) is constant, hence
MMDD)((P)) Faj1, (Po) F
(34)
S c(0) (MDD((P))#1,, (Po)+u) for any unit Φ,
(35)
so the desired bound holds trivially.
Case 2: |col|| > O. Write ro = =I[wol| and Φ0 == wo/ro E gd-1 For any X ~ and Y ~ u,
(Po o (X) = (0, AX + b) (0,b) (0o. X),
and similarly for Y. By (T) and (S), obtain
(MD)((P)) F+ (Po) MMDs(Law X)),
(36)
c(ro) MMDs((Pa.. ( # (Po))#V) c(ra)
(37)
sup MMDD((PP)HH) (Po) $E84-1
(38)
Since ra-
||A 0ll < ||AT|l ||Allop L and L [0, 1]. have ro [0, 1] that c(re) is
well defined.
Taking the supremum. Now take the supremum over 0 gd-1;
sup MMDxMMDi((E (Po)MMDDDPooPPPPP (Po) I sup c(re) sup MMD( (P))
(39)
Using re S and the fact that nondecreasing on [0, 1]. we obtain
sup (Po) F414, (Po) F S c(L) (SMMDA((,.)
(40)
The left-hand side is exactly AFMMMD((#HH. F41), which proves the claim.
Lemma 14 (Max-slicee sliced mixture p-covvexity for MMD,). Let k RxR + R be a kernel on R. and let MMDj be its associated maximum mean discrepancy P(RR). Assume that MMDs is mixture p-convex for some p E [1, bo): for every probability space (Ω, F. F,P)
(Pe), (ve) C P(R)
for which = So Pe p(de) and D = Sq Ve p(dc) well defined,
MMDR (L. p(de), I Ve peplle) <
(J, PM p(de)
1/p
Define the max sliced lift on ((Rd) by
(H, : sup 0E3E-1
(MD) (Po)
Pa(x) = (f,
Then MSMMD, is also mixture p-convex:
MSMMD ( Ihe p(de), I ve p(de)) s
511 MSMMDe u.) p(dc)
1/p
42

---

Published as a conference paper at ICLR 2026
Proof. Fix 0 sd-1 and set By mixture p-convexity of MMDs on IR,
Me (oo)
== (Po) E EP(B).
Pushforward commutes with mixtures:
(Po) pn dp) - we dp,
(Po) (pd = ue dp.
(MMo Pe dp, (v))= $ (ʃ (/MMD dp)
(41)
Taking the supremum over A on the left-hand side of Equation 41 gives
sup MMD; ((Pb) # Me dp, (Po)# S vo dp) s sup (ʃ 2M
v2) (42)
Define f(0,c) MM vn) and h(c) sup, f(0,c) &SMMD Since
f(0,c) h(c) pointwise in c, obtain for every A.
( ((.copppoc) 1/0
Taking sup, yields
1/2 sup (M < *( (MMSMMDA(We.. (He: u2) dp) )/p (43)
Combining Equation 42 and Equation 43 shows
MSMMDA (HHole Ho... dp. I [vo dp) s (/ {/MSMMDD&W v.) p(de)
1/r
as claimed.
Theorem (Supremum-max mappemmmm-aax-siieed MMDs contraction of the complete Sobolev Bellman opera- tor). Let
T (S x A P(R
(s x A P(
be the complete Sobolev Bellman operator bootstrupping the vector (2,0,2,0,7)
Let l: R TR + R be kernel TR. and let MMDs be its associated maximum mean discrepancy on P(R).. Define the max-sliced lift on P(RRI+m+R) by
HSMMDD((HR)) = sup
MMD (PPo)((())))) (Po)&V),
Po(x) = (0, x).z).
and its supremua version over state-action pairs by
(mM == sup MSMMD, n2(s (s.a)
Assume MMDs satisfies:
Translation invariance: for all 1 R and all μ, P(RR)
MMD (MDD&((++>) (x+ +++++t))++++ (x+ x t) MMD (H, v).
(S) Seale-contraction there exists a nondecreasing c [0, [0, oo) such that for all [0, 1] and all p.v E P(B).
MMDx #1) (xr sT) S M())
43

---

Published as a conference paper at ICLR 2026
(Mp) Mixture p-convexity: for some p [1, oo). for every probability space (S, P) and families (He), (ve) C P(R. for which the mixtures are well defined,
MMDE (L. (/uaplddc) I ( S <(I (JMMMDAUUOs uc) p(de)
1/p
The complete Sobolev Bellman operator is defined by the same reparameterized ingredients as in Theorem 3 namely the policy (s) the transition f(s, (ss,a;Es): and the reward r(s, In particular, ar, for each (8, u) and noise draw (e,,Ex) the complete Sobolev update acts as an affine push-forward with linear part L(8, 0, (, (as in Theorem Define the uniform bound
|| = sup sup [(s,a;E s,a Ep
sup
S Ty Ktull-
Assume ||al E [0, 1]. Then for any two Sobolev return couum--iisribbttiin functions Th 722-
(SMMDI(TT T))) S c(l())a) MSMMDA(1, n2).
In particular, T S is a strict contraction whenever cII) < 1.
Proof. Fix (s, a) and condition on C' = (Er, Set r= r(s, ==r(s,a;er) 8' f(s,ae)))
d' =<s;ee)))
and let X Equation 22
C (s), By the push-forward identity (Lemma [ and the affine update in
Law (x)))x.(X)))
Affine push-forward fixed Let L(C)- := [| a; y, E &x) )lop. Applying Lemma 13 (for MMDs), which relies on (T) and (S), to the conditional laws of X gives
MSMMD; (Law(Q,.aXXi) (), Law( (X2) (O)
S <e(L(C)) MSMMD (La C), Law(X; C)) = =c([(C)) MSMMD (71 (s',,'),,
(44)
Averaging over C. [4 MMD,),, which relies on (M,). together with Equation yields La
MSMMD,
1/P
(IM C)2, Law @ (X,) CJ)" p(dc))
(45)
1/p
(c(T(C))
p(dC)
(46)
1/p
<o((llaaa) (MMISMID. (M(ss (s', ad))) p(dC)
since is nondecreasing on [0, 1] and L(C)= < [[]]| < for all C. Supremum bound. For any realization of C, by definition of the supremum metric, MSMMD& (m (6',a'),,, aaa), n2(5), ,d)) s sup MSMMD, (m (v, v), re(u, v))
(u,0)
MSMMD (h, 12).
Combining this with the previous inequality,
MSMMD, ((TS mem (s, a). (2),aa)))
-SMMMD, (Law Law (1 (Xx)))
1/p
S sc(l) (MSMMDD.(...)) p(dC)
=((1 MSMMD (p))
44

---

Published as a conference paper at ICLR 2026
since MSMMD n2)) does not depend C. Taking the supremum over (s, a) completes the proof:
MSMMD (T (TS.. m1
T ne) s dlll MSMMD) 12)-
Corollary 5.1 (Supremum- pappremum-mx--aeee max-mmma--------eaaaa slacsl contraction for MMD with the multiquadric kemel). Let
(x, y) -VI+ 1-_ 2 ylI with h > 0 (the multiquadric kernel). Consider the setting of Theorem 5 and assume in addition that the bound ||L satisfies [| E [0, Then, for all
return cruum-dittbbution functions 1/1, 72,
MSMMD k. (T (T
m. T I VlllIa MSMMD (n1 n2)-
Proof. We verify the conditions of Theorem for the MMD associated with the multiquadric kernel.
(T) (Translation invariance): Since k, radial, depends only on the distance ||a: yll- implying that the MMD induced by ky is invariant under simultancous translation of its arguments.
(S) (Scale-contraction): By Lemma 10 (specifically Equation 28 for any 8 E [0, 1]. the unsquared MMD under kn satisfies contraoonta with the specific function c(s) v8.
(M,) (Mixture p-convexity); The kernel kn conditionally positive definite (CPD) (see e.g., Theorem 3.1 of Killingberg & Langseth (202)). Therefore, by Lemma 12 the MMD
associated with kn, satisfies mixture p-convexity.
Applying Theorem with c(s) v yields the contraction factor e(IlC]la) VTClla-
Theorem 6 (Supremum max sliced MMD contraction of the incomplete Sobolev Bellman opera- tor). Let T (AxA-) ) Sobolev Bellman operator that updates H (Z, Zz in Equation 18
+ (S x A + P()!+T)) be the action graddent (incomplete)
Assume:
Translation invariance: MMDa((x + () x () MMD (1, v) for all L E R and all p. E P(R).
(S) Scale-contraction: There exists nondecreasing [0, + [0, oo) such that
(MD + 41) (x + sz) S (*(s) MM V [0. H,V P(R).
(Mp) Mixture 1-convexity: For every probubility space (9,F,,) p) and families Hey Ve P(R) for which p == Sa Me p(dc) and = Sa Ve p(de) are well defined,
MMDR (A I _ papllc)
MMDx ve) p(dc),
and the corresponding max-sliced version (Lemmal 14 holds.
Assume
llfall2 S Lsas
||Talle S L.
For the linear blocks in Equation 18
A(s,0)E/.
=x)
N(s, =
fo T
Define
:= sup 8,0,8f.8m
[[[A N N]ll
and assume L, S 1.
For each (s',al), assume there exists lifting map
Js p(R)+m)
((R++++)),
45

---

Published as a conference paper at ICLR 2026
with first marginal equal the input law and joint law induced by our distribution parametrization (Appendix B.2), so that the second marginal matches the true law of O,Z'. Assume that for some Lang 2 L.
MSMMDA (Tss ( Jo S Laug MSMMDA((4,), Vp,u.
Define MSMMD 72) = (a.() MSMMDD((M(ss.)) a), (2(s, maps Th 72.
Then, for all return-gradien
MSMMDD((T?2)) T2) Lang (x(L) MSMMDA 72)-
In particular, if Laug eu (L) < 1, then T is strict contraction.
For the multiquadric kermel ky (2. y) = - - 12 yll?, Lemma 10 gives cx(s) v8 for E [0. Hence, whenever T. S 1, sufficient condition for strict contraction is Trang VT. < 1.
Proof. Fix (s, a) and a noise draw C = EE) Set
= f(,a;EE))
a' = ='s:EE))
We will rewrite the incomplete update affine map augmented next-step variable that includes the missing state-gradiens component.
Augmented variable. Let X (H), U!) E +1+m) with law ne (n..a Define the affine map on the augmented space
s.a ( C) b + NJE.
Then, conditionally on C, we have
(0, as (X: C)) (0, (IA NITO, X!).
Moreover, by the same argument as in Lemma
(T a) Law T s,a (X(CC)
Apply max-sliced affine contraction lemma. Since T. < 1, the operator norm of the joint matrix ||[ N] [0, 1] for all C. We may thus apply Lemma 13 (which relies on (T) and (S)):
MSMMDA (La C)) C), Law( (X,: (X); (X
(IlA ]II) (f,, 72)
Lift back to (Z, 8). By the lifting assumption,
MSMMD (fn. < Laug MSMMD m2(s'sd'))
Average over noise. By Lemma 14 (mixture 1-convexity), have
MSMMD ( m (8, (T (s,
MSMMD, (Law (x1:C) C), Law (Xt;; ()) p(dC).
Combining the affine and lifting bounds derived above, the integrand is bounded pointwise by Lang (Ill A N]II) MSMMD (s', a'), n2(s', .a))
By definition of the supremum metric, MSMMDD (u, v), (2 v)) MSMMD (m, 922 for any state-action pair (u, v). Using this global bound and A N < L, the integral satisfies
MSMMD, S (s, a), (TS")))) a)) < Laug (L) MSMMD (1h, 12).
Taking the supremum over (s, a) the lent-hand side yields the claim.
46

---

Published as a conference paper at ICLR 2026
H BACKGROUND ON THE WORLD-MODEL
H.1 CONDITIONAL VARIATIONAL AUTO-ENCODERS
A principled invertible generative model can be obtained from Variational
(VAE) Kingma 2013). More interestingly for us are conditional VAE Soh et al. (0013) which we briefly introduce.
Conditional Variational Autoencoder (cVAE) is a generative model that learns to generate new samples from distribution conditioned given input information. In case, the cVAE models the distribution of next states and rewards conditioned on current states and actions.
Formally, the vVE consists of components:
Encoder:
The encoder qc(E s', s; a) maps the observed next state s' and reward r, conditioned on the current state-action pair (s, a), to latent variable e, typically modeled as Gaussian distribution with diagonal covariance matrix:
(e(5 sr; ,, = NNs (1..
(47)
Decoder: The decoder Pel (s',r s, a) reconstructs the next state s' and reward from the latent variable E, conditioned on the current state-action pair (s, a).
Prior: The prior Pa (s', 8; a) allows more flexibility in the latent space than simple standard Gaussian N N(;; T) It is also parametrized as a multivariate diagonal Gaussian:
Po(e 8, a, (l, oz
(48)
The objective of a eVAE is to maximize the Evidence Lower Bound (ELBO), which balances accu- rate reconstruction of the input with regularization term that ensures the learned posterior distribu- tion remains close to the prior distribution. The objective as follows
LeVAR (S, vb) Bygcelle" [log (s', e3s
(49)
-ÀKL x DKL (qc(E s,,r:s,a) || Po(E s,))).
The first term encourages faithful reconstruction of the next state and reward, while the second term regularizes the posterior distribution remain close standard Gaussian prior. Assuming the decoder (s(s', e; 8, a) is Gaussian with a fixed variance, the reconstruction term reduces to an 1.2 loss, which can be estimated using the difference between the reconstructed samples and the true samples.
Assuming the encoder parametrizes Gaussian with diagonal covariance and that the prior is also Gaussian with identity covariance and zero mean, the KL divergence can be estimated from encoded input samples as
d (1 + log(o (s, - 2 as, (0(0
(KL(C) = (s,a)))
(50)
2
H.2 CONDITIONAL NORMALIZING FLOWS AS WORLD MODELS
Normalizing flows provide an alternative class of generative models based on a sequence of invert- ible transformations with tractable Jacobians. Let fo : Rd - IRd be invertible map, and let z ~ Pz denote a simple base distribution (e.g. standard Gaussian). A flow represents a random variable through the transformation fo (z). Its density follows from the change-of-variables formula:
Px (x pz Pzw (z)(z) |det (x)(
(51)
By designing each transformation so that its Jacobian determinant is efficient to compute, normaliz- ing flows permit exact likelihood evaluation and fast sampling.
47

---

Published as a conference paper at ICLR 2026
RealNVP. RealNVP (Dinh et al. 2016) is widely used normalizing- flow architecture based on affine coupling layers. Each coupling layer partitions its input (x1, 12), leaves I1 unchanged, and transforms 22
= 02 @ exp((s<(x1) + (21).
y1 = 21,
where 86 and &s are neural networks. Because the transformation of xg depends only on T1, the Jacobian is triangular, and the log deg-determinant reduces sum the scaling terms. Stacking multiple coupling layers with alternating partitions yields an expressive model with computationally cheap density evaluation and sampling.
Conditional flows for one-step dynamics. To model one-step transitions, the flow is conditioned on the current acte-actionn pair (s, a). The coupling-layer networks sx(-) and t-)-) therefore take (s, a) as an additional input, yielding a conditional distribution
Po(s', s,a)..
Sampling proceeds by drawing z ~ pz and applying the inverse transformation
(s',r) fo (z;
which provides differentiable mapping from (s, 0, 2) to next-state and reward samples. As with the cVAE model introduced above, this enables efficient sampling together with the pathwise derivatives ((s', r)/0(s a) required for Sobolev temporal differences.
I ALGORITHM DESIGN, PSEUDO-CODE AND BASELINE
In this section we describe the motivation for some design choices of our method. Section [.] sets the design for the toy reinforcement learning experiment of Section 6. and the MuloCo experiments
from Section 6.2 We follow by providing the pseudo-code in Section 1.2] Finally in Section I.3 we
discuss our implementation of our primary baseline MAGE D'Oro & Jaskowski (2020).
I.1 DESIGN CHOICES
Firstly, most experiments are conducted in the data-efficient setting, where the number of updates per interaction with the environment (UTD ratio) is larger than one, making stability and overestimation bias critical concerns. MAGE adds gradient regularization which itself incorporates several modifications to DDPG (Lillicrap et al. 2016). Below, we describe each modification and whether it was retained in our implementation.
top of TD3 (Fujimoto et al., 2018),
Target policy smoothing: this technique applies random noise to the policy's 's action when esti- mating the TD learning target. It the only TD3 (and MAGE) modification we retained, other changes cluttered the implementation and introduced noise that could interact poorly with distribu- tional modeling. It is worth noting that omitting target policy smoothing in similar settings is not uncommon (Singh et al. 2022 Kuznetsov et al. 2020).
Delayed policy update: the policy updated once for every two critic updates, which helps stabilize learning by preventing premature policy shifts.
Double estimation: an ensemble of two critics, Q1. Q2 or Z1. Zz, is used, with the bootstrapped target taken as the minimum of the two estimates at the next state (including policy smoothing noise). This modification introduces an underestimation bias counteract the well-known overestimation issue in value-hased methods (Van Hasselt et al., 2016). We observed double estimation to be critical for stable performance in the data-efficient setting: without it, average Q-values quickly diverged. Overestimation bias can also be addressed in the distributional setting, as demonstrated by TQC (Kuznetsov et al. 2020). A straightforward approach to induce underestimation is to truncate the top p% of values the critic's target distribution; in our setup, truncating few as 25% of values proved highly effective. Following TQC (Chen et al., 2021), we employ an ensemble two distributional critics whose samples concatenated.
48

---

Published as a conference paper at ICLR 2026
1.2 PSEUDO-CODES
In this subsection, we present the pseudo-code for the MSMMD estimation procedure, the MMD estimation procedure, and the full algorithm for policy evaluation and improvement.
(Estimation of MSMMD); approximates the supremum over projection di- rections by gradient ascent on the unit sphere, yielding an empirical max-sliced MMD between two sets of samples.
Algorithm I
( Estimation of Sobolev samples): leverages the world model to generate differentiable Sobolev-return samples, and applies truncation to mitigate overestimation bias as Kuznetsov al. (2020).
Algorithm 3 (Full DSDPG algorithm): integrates Sobolev-MMD components into the DDPG framework Lillicrap et al. (2016).
Algorithm Estimation of MSMMD from empirical samples
Require:
Empirical samples X ( Rd. Y
N R2
2: Require: Kernel i defining the MMD
3: Require: Number of gradient steps T, step size y
4: Output: MSMMD(X,Y)
5: Draw w (a)
(random unit direction) frandom unit direction
6: 0 <- w/lk
7: for 1,.... T do
8: for ...... N do
9:
Ui -- x;))
{project to ID along 0 {project to 1D along 0}
10: Ui (0, y1)
11:
end for
12:
J(Ø)= + MMDD((u),, (u)N))
maximize MMD over 0)
13:
g (0)
{gradient w.r.t. direction)
14:
w w +
(sccent step in unconstrained space). (re-normalize onto the unit sphere)
15:
0 w/u/lu
16: end for
17: 8 + -toop stop_grad(())
(stop gradient the final direction)
18: Return: MSMMD(X,Y) + (MD,(((0, {00, 9.) )
49

---

Published as a conference paper at ICLR 2026
Algorithm 2 Estimation of MMD? loss via imagination of transition samples with Sobolev samples
1:
Require: Number of samples M, kernel k. discount factor y E (0,1)
2: Require: Truncation percentage E [0, 100]
3: Require: Distributional critic (s 0,
4: Require: Policy network xo (s))
5: Require: Conditional VAE (cVAE) with prior (e(e) and decoder (s 8, 8,),e)
6: Input: Transition sample (s,
7: Input: Online critic parameter Φ, target critic parameters of and 062
8: Input: Target policy parameter f'
9: Input: Boolean flag use act ion gradient
Output: Gradient estimation of MMD with respect to
11: Draw (E)
{Sample latent variable from the prior) (Generate transition using the decoder {Action from target policy on s') (Samples from online critic)
, (p 13: a' + (3))
(8',F)~Pu(
6)
14: Sample ZI:M
i.i.d. Z (8. a)
~
15: Sample _ 1:M
i.d.
Zoz ,d),
Sample 2.22 I:M
1e
i.d. (3g(3',a' d')
(Samples from each target critic} [Concatenate target-critic samples)
1:x concat Z 2 18: if uuseacctio.ggaiieer then for each 15 2M do
20: Y f.,.1.(( (?neut.)
(Bellman target Sobolev samples) { (Action gradient of Bellman target)
21:
V.Y; + Vali. (Znexta)
22: Y concat(Y, VaYi)
23: end for
for each < < M do
25: V.Z. + V az(
(Action gradient of online critic)
26: 2: + concat Zi Vali)
27: end for
28: else
29: for each < i < 2M do
30: Y, - PS ,, (Zeext
31: end for
32: end if
33: Prune top 9% of (
12M by return magnitude 34: Compute 2MD
[Low-bias target set)
35: MMD E k(Z,, 2z) 2k(2., Y2) + k(Y,
Return: MMD
50

---

Published as a conference paper at ICLR 2026
Algorithm 3 Distributional Sobolev Deterministic Policy Gradient (DSDPG)
1: Require: Number of samples M, number of policy samples Mpolicy, kernel k discount factor
7 E (0, 1). learning rates ag; Ot.
2: Require: Two distributional critics Z (s, a) and (s, with targets Zoz (s. Zog (s. a)
3: Require: Actor network (s)s) with target (s)
4: Require: Replay buffer '
5: Require: Boolean flag use lae.aattii act.ii.nggriiient
6: Require: Parameters cVAE: Encoder ((). Decoder (1²), Prior (V)
7: Input: Initial parameters 0, Φ1, d02 and target parameters 0' 0, 01 t- 21, + 02
8: Input: Policy update frequency
9: for each episode do
10: Initialize a random process M for action exploration
Observe initial state So
12: for each step in the episode do
at f (sx) N,
Seelect action with exploration noise)
14:
Execute as, observe (,
Store (St, Tt, Tt, 81++1 in replay buffer 5, train_world. modell St, ae, sc+1,r))
16:
(Train eVAE world model and update
parameters)
17:
end for
18:
for each gradient step do
19:
Sample mini-batch ((s;, D Compute Distributional Loss via MMD (cf. 2):
(Replay buffer sampling)
20:
NB. Samples are concatenated and top p% are removed TQC Kuznetsov al. 2020
Lz, MMD Sobolev (Zo, (Zo, Zoz 84, as To Y, use action www.aaictnnniniiieee gra.dend))n
Izz MMD Sobolev (Zoz. (Zoz Zoz,) ); 8i, as moy Y use acti on gradient)
24:
Update critics:
Φ Φ1 0s Ve, Lz
26:
62 t as V Φ2 Lz,
27:
gradient step mod d then
28:
Sample Mpotiey values from critic for each state in the batch:
(6)
iid. Zo. (8t, To(s)) for each i f1...., N}
29:
"policy.1-Mpetz
30:
Compute actor loss using double summation:
L2 = N Σ Malky E z
31:
p
32:
Update actor: 8 &- ag VoL
33:
Update target networks:
34:
TΦi + (1 )
35:
0 - TØ2 + (1 7 o
36:
0 - r0 (1 7)0'
37:
end if
38:
end for
39:
end for
I.3 BASELINE MAGE
The primary baseline of this work is MAGE D'Oro & Jaskowski 2020). they the first to propose to use the action gradient to steer policy evaluation. The primary difference with our work they did so deterministically. Hence MMD setup collapses to simple regression setting.
51

---

Published as a conference paper at ICLR 2026
As primary baseline, consider the deterministic counterpart of the Soboley-distriuutional backup. Define
Sdet (s, a) del (s, a)]
fdec(s, a) =
Sdet (s, a) = r(s, + (
fdet (s, a) Vor(s,a) +1 VoQols',
L Lo (Ø; s, a, ) = | (s, (s,a) ao(s,a) + As f|lc (s, a) Volo(s.a)
In practice, s' and r are sampled from the stochastic world model (cVAE) as in DSDPG. Further- more, the above L2 terms are replaced by Huber losses, with the gradient-term weight set to 5 just as in Oro & Jaskowski 2020).
ToY SUPERVISED LEARNING
To motivate our algorithm, demonstrate its ability to learn the joint distribution over both the output and gradient of a random function supervised setup. compare deterministic Sobolev training Czarnecki al. (2017) against Distributional Sobolev training and show that only the latter can capture the full variability of both outputs and gradients.
E200
(a)
(b)
(c)
Figure 4: (a) Samples of the marginals of the full Sobolev distribution [f(æ; a); ( a)]: output (left) and gradient (right). Blue: samples from the true distribution; red: samples from the Distri- butional Sobolev model trained via MMD; green: samples from the deterministic Sobolev baseline Czarnecki et al. (2017). (b) Biased MMD score (lower is better) on the joint variable. (c) Lg error between predicted sample mean and true mean. NB. Predicted (red and green) and true (blue) samples are highly overlapping.
The task involves learning one-dimensional conditional distribution p(y[z), defined as mixture of sinusoids f(x; a) = a x sin(x), where the latent variable a is uniformly drawn from {0, 1, 2, 3, 4). The distributions over outputs f2; a) and their gradient Vaf(2; a) depicted in Figure 44 It compares an MMD-based model and a regression-basse model, trained using stochastic gradient descent with identical architectures. In unlimited data regime, new pairs of and for each batch, and for each x, four a values were sampled with replacement, yielding samples (xr 31.1).
drawn
As expected, the MMD-based model captures the full joint distribution [f(I;a); Vaf(a; a)], whereas the regression baseline collapses to the conditional mean [f Va) v.I(z;a) Fig- ure 4b plots the MMD score (lower is better) on this joint variable using an evaluation kernel, and Figure 4c shows the L2 error between the regression model and the sample mean of the Sobolev generator. As result, the MMD model better matches the entire distribution, the regression model slightly outperforms on the mean, and both methods effectively exploit the gradient signal (blue curves).
The distributional variant via MMD and the regression one via L2 used the same architecture except for some noise of dimension 10 drawn from N /O; /) concatenated to the input for the distributional generator. For each pair (xx 3/44 ). four samples were drawn from the generator. Both were trained using Rectified Adam Liu et al. (2019) optimizer with learning rate of 1 x 10-3 and (3o. 31) (0.5, 0.9). Neural network is a simple MLP with 2 hidden layers 256 neurons and Swish non- linearities (Ramachandran al. 2017).
52

---

Published as a conference paper at ICLR 2026
Maximum Mean Discrepancy (MMD) estimated using mixture of RBF kernels with band-
widths o; from the set 01 02...., 07} = {0.01, 0.05, 0.1, 1, 10, 100). We used the biased
estimator from Equation 25
The equation for mixture of RBF kernels is given by:
|| y
k-mix (xr, v)
exp
(52)
20
The evaluation kernel we used was the Rational Quadratic A.QQ with with
RR (z. v) (1 1
| u1l
(53)
2a
Regarding the dataset, the (I,, (1:4) pairs were drawn with U(0; and was drawn from (0, 1,3,3 4] with replacement. In the limited data regime, the pairs (Ir, V1:4) were drawn once and stayed fixed. The batch size was thus equal to the number of points in the dataset. In the unlimited data regime 256 new pairs were drawn for each batch. Every experiment was run for 25 000 batch samples and thus the same number of SGD steps.
Limited data regime In both supervised and reinforcement learning tasks, the assumption of un- limited data is unrealistic. Here, explore how the performance of the MMD- based and regression- based models diverges when the amount of available data is restricted. We use the same setup as before, but with fixed number of (ar, 31:4) pairs. Several aspects of the learned models can be inspected. In order to assess stability, we report the average norm of the second-order derivative over the input space. For accuracy, we measure the average L2 losses between the true expected gradient and the predicted gradient. Results are shown in Figure [S] As can be seen, the deterministic model tends overfit rapidly, while the distributional variant (MMD) proves more robust, main- taining better performance even with constrained data. Notably, the recond-orde derivative for the deterministic model escalates sharply as data becomes constrained, indicating instability in its ap- proximation. In Appendix J.1 we discuss different common tricks to mitigate overfitting and related issues in RL.
distance [V,(t) with limited data
Average of derivalive
10'
Hessien rorm (lapccale)
Legend Detenministic Stochestic
distaance (looscale)
Sobolev
109
Distributtonal
3.0
Deterministic
10-3
3
512
1024
untimited
10
15
20
25
Epach
Number of training pains
(a)
(b)
Figure 5: Toy supervised learning problem. Comparison between MMD-based 1.2 based model- ing. Left panel: training curve of 1.2 loss (logscale) on gradient between true conditional expectation with regression prediction and with empirical mean from MMD-based model. Sobolev (blue) used gradient information to train either using MMD (full line) or L2 regression (dashed linc). Right panel: average over the input space of the second order derivative (logscale) of predicted gradient from deterministic model (blue), MMD/ / stochastic (yellow) and with gradient information Sobolev (dashed). Metrics averaged over seeds.
J.1 COMMON TRICKS
Overfitting with limited data a common issue in regression tasks. Early stopping seems an obvious solution in this case but emphasize that it requires evaluation criterion that is not always
53

---

Published as a conference paper at ICLR 2026
available (for example, in policy evaluation). Other solutions include weight regularization Krogh & Hertz [1991) dropout Srivastava et al. (2014), Bayesian neural networks Blundell et al. 2015), ensembling Chua et al. (2018), and spectral normalization Zheng et al. (2023), all of which often reduce network capacity. To address similar issues, Fujimoto et al. (0118) proposed adding noise to the target match, effectively smoothing the critie. As argued by Ball & Roberts (202T). this method can be seen as indirectly acting like spectral normalization, encouraging smoother gradients and effectively reducing the magnitude of the second derivative. Appendix D.2inddxxx..jooo shows how noise scale impacts overfitting by inducing bias. On the other hand, we propose avoiding such assumptions by using generative modeling add latent freedom.
J.2 ADDING NOISE
added some independent noise on for each (x,, V1:4). Noise scale Ø was in f0.01, 0.I, 0.53. For each new batch, it was sampled from standard Gaussian " NO o²) and added
Inspired by Fujimoto al. 2018
x +1
In Figure 6 7 we can see the impact of the various noise scales on the predictions of the deterministic regression. As can be seen, adding noise on has positive effect in terms of stabilizing the gradient but it induces a bias that grows with the scale of the noise. Moreover, this noise depends on the application and makes strong assumptions about the function learn. The stabilizing effect of additive noise can further be seen in Figure g where both the L2 loss and average second order derivative are displayed as a function of the number of sampled locations.
True and predicted f(x)
True and predicted v.f(x)
True
No Nolse
Ouipul values
Grad ent volues
-1-
-2
-3
values
Figure 6: Toy supervised learning problem. Comparison of samples from the true five-mode dis- tribution with predictions made by deterministic model trained with L2 loss (green). The output space is shown on the left, and the gradient space on the right. Results obtained after 25,000 training steps.
True and predicted
True and predicted 7,1xx
True Noise 0.01
Noise 0.1
Noise 0.5
Gradient values
Output
-6
2
3
values
valuns
Toy supervised learning problem. Comparison between true samples from the distribution of five modes and deterministic models trained with varying levels of additive noise on their input data. Low level of noise (green), medium level of noise (orange), high level of noise (red). Results obtained after 25,000 training steps.
Figure 7:
54

---

Published as a conference paper at ICLR 2026
Average L2 loss on V,fxx
Legend
Noise scale: 0.01
L2 mean gradients (logscale)
10a
Naise scale: 0.1
Sobalev
10
100
10
3
128
512
1024
unlimited
training pairs
Figure 8: Toy supervised learning problem. Comparison of the L,2 loss between the predicted gra- dient and the conditional expectation of the distribution. Different scales of additive noise on the input are compared: low noise (light blue), medium noise (medium blue), and high noise (dark blue), alongside Sobolev training (dashed). Results are shown after 25,000 training steps.
Average norm second order derivative 72f(x)
Legend
10a
Neise scale: 0.01
Neise scale: 0.1
Hessian norm (ogscclle
Noise scale: 0.5
Sobolev
101
109
-3
128
532
1024
uniimited
Number of training pairs
Figure 9: Toy supervised learning problem. Comparison of the average second order derivative norm over the input space. Different scales of additive noise on the input are compared: low noise (light blue), medium noise (medium blue), and high noise (dark blue), alongside Sobolev training (dashed).
55

---

Published as a conference paper at ICLR 2026
K TOY REINTORCEMENT LEARNING
K.1 ENVIRONMENT
In this appendix, we provide full details of our custom toy environment, summarized Table 1
Environment Dynamics. A point mass moves in two-dimensional continuous state space, and the agent controls its acceleration in [-1,1] per axis (Acceleration range). At each step the mass's velocity is updated via Euler integration with friction (Friction coefficient) and then its position advanced by the new velocity (Integration time step). The mass has unit mass (Mass of agent) and radius 0.5 (Agent radius), and if its center leaves the square of half-width (Bounding box half-width) the episode ends with no reward (Reward for leaving bounding box).
Partial Observability via Memory. At the start each episode the mass is placed uniformly in square of half-width (Initialization area half-width). We then sample one of M E {3, 4, 5. 6} hidden bonus locations arranged radially around the center. Reaching the correct location of radius 0.5 (Bonus radius) yields terminal reward (Reward for reaching bonus). The agent does not know which location is active, but each visit to location sets a binary memory flag, making the MDP partially observable and forcing exploration.
Controlling Distributional Modes via M. By sweeping M E (3, 4,5, 6} we directly tune the number of modes in the return distribution across episodes, from concentrated when M is small to highly multimodal when M is large.
Parameter
Value
Maximum episode length
25 steps
Reward for reaching bonus
10
Reward for leaving bounding box
0
Acceleration range
[
1] per axis
Mass of agent
1
Agent radius
0.5
Bonus radius
0.5
Bounding box half-width
3
Initialization area width
1
Friction coefficient
0.1
Integration time step
0.5
Table 1: Summary of the toy environment's key parameters.
K.2 EXPERIMENTAL DETAILS
The settings were similar to Section apart from the TQC Kuznetsov et al. (2020) truncation pa- rameter p which was set to 5%.
L REINFORCEMENT LEARNING EXPERIMENTS
Here we describe the architectures, optimizers and other hyperparameters of the full Distributional Sobolev Deterministic Policy Gradient algorithm.
Policy network Policy network is MLP with hidden layers of 400 neurons. The non-linearity was Swish (Ramachandran et al.) 2017). Final activations are mapped to the output space using linear transformation followed by a tanh non-linen non-lincarity. The policy network was optimized using the Rectified Adam Liu et al. (2019) with learning rate of 1 x 10
Critic network Critic network architecture is almost the same for the policy network. Im- portantly, it is kept constant for experiments using normal DDPG and DSDPG apart from noise concatenated the input (s; a). The network is MLP with 400 neurons, skip connections from the
56

---

Published as a conference paper at ICLR 2026
Table Hyperparameters for the DDPG and DSDPG experiments on MuJoCo environments
Item
Value
Discount 7
0.99
Polyak averaging
0.005
Buffer size
10
Batch size
256
Exploration noise scale
0.1
Critic learning rate
1 10
Policy learning rate
1 x 10
cVAE learning rate
1 x 10
cVAE weight decay
1 x 10
cVAB KL weight
0.1
cVAE latent dim
S +
Critic input noise dim
64
Number of samples dist.
10
%p truncation (TQC)
25%
maximum slicing optimization steps
100
maximum slicing LR
le-4
maximum slicing optimizer
Adam (Kingma 2014)
input and Swish activations. No non-linearity applied on the output after the last linear layer. The critic network is optimized using Rectified Adam Liu et al. (2119) with (31,32) = (0.9, 0.999) and learning rate of 10-4
Conditional VAE world model The encoder, decoder and prior networks are MLPs with 3 hidden layers, each containing 1024 neurons. Skip connections are applied from the input to each hidden layer. The cVAE is optimized using RAdam with weight decay 1 x 10-1, (2) (0.9, 0.999), and learning rate x 10-4
The prior is learned diagonal multivariate Gaussian N(= (s. a), o a) C 1) with latent dimension equal to the size the random variable being modeled, which is |S] 1 for (s', Following D'Oro & Jaskowski (2020); (huu al. (2024), the cVAE predicts the difference between the current and next observation, d, = which is then added back to s, along with the reward
Conditional Generative Moment Matching For the distributional critic, noise vectors were con- catenated with the state-action pairs ()) and passed through the same architecture as the deter- ministic critic. The noise dimension was set to 64. Noise was transformed by an independent 2 layers small MLP width 64 with Swish activation before being passed to the critic. For cach state-action pair, 10 samples were drawn to update both the critic and the policy. The multiquadric kernel (2, 3) -V1 + b2ll y Killingberg & Langseth (2023) was used, with the kernel parameter It set to 100.
L. FULL CURVES AND WALL-CLOCK TIME
The full evaluation curves on the 6 environments are reported in Figure 10 The wall-clock time of the different methods is displayed in Table
M ADDITIONAL EXPERIMENTS -RL
In this section provide further experiments along three axes: kernel bandwidth (Section M. noise scale (Section M.4). number of samples in the distributional methods (Section M.5) and ca- pacity of the world model (Section M.6).
57

---

Published as a conference paper at ICLR 2026
Marmal
Muttipiiative
Gavesian naksa
12:000
11:00
100.00
21:00
29.00
Ne
1at
100
27
100
128
150
25
200 800
128 1000000
50.000
20.000
85000
50.000
250
H50
20
and
op
50
20.00
100
0.50
2000
125
150
2..00
000
150 000
100 S0
Raa
Side
2000
1000
100
(TD3)
MND
NSVMD MSVND Sobalev
(MAGE)
MND Sobalev
Evaluation of DSDPG (MMD Sobolev), deterministic Sobolev/MAGE D'Oro & Jaskowski 2020), TD3-Huber Fujimoto et al. (018). IQN Dabney al. (011a)). and standard MMD Nguyen et al. (2020); Killingberg & Langseth (2023) six MuJoCo tasks. Results are re- ported over 10 random seeds. The median is displayed with 25%-75% IQR. Window smoothing with window size 3.
Figure 10:
58

---

Published as a conference paper at ICLR 2026
Method
Time for 1000 iterations (s)
MSMMD Sobolev
62.5
MSMMD
45.5
MMD Sobolev
40.0
MMD
35.7
IQN
35.7
Huber Sobolev
31.3
Huber
30.3
Table 3:
Wall-clock time perform 1000 iterations of each method under Humanoid-v2. Experi- ments were run on a single Nvidia H100 GPU.
M.1 CONDITIONAL FLOW WORLD MODEL
The conditional flow world model is implemented using RealNVP architecture (Dinh et al. 2016) Each affine coupling layer parameterized by conditioner MLP with two hidden layers of width 512, using SiLU activations (Elfwing et al. (2018) and LayerNorm (Ba et al. (2166) after each hidden layer. The linear transformations within the coupling blocks use the invertible 1x1 transformation parameterized through an LU decomposition which enables learned invertible mixing of feature dimensions while keeping the log-determinant of the Jacobian inexpensive compute. As in the cVAE world model, the flow predicts the difference 6, - along with the reward r, and the next state reconstructed as s + 8,.
introduced in Glow (Kingma & Dhariwal (018)).
Following the same output parameterization the cVAE world model, the flow predicts the con- catenated vector (6,, r) of dimension [S] 1, where 8₂ = s' s. The next state is then recovered as s' 6,.
Training proceeds by maximizing the conditional log- likeliiodd log (o., (r..r 8, a) using the same optimizer as for the cVAE. Sampling draws ~N(0, from the base distribution and applies the inverse flow conditioned on a).
(6,,r) = (
Because the model is fully reparameterized, each sample (5..r) admits pathwise derivatives r(),,r) a), which are required for computing Sobolev temporal-differenn targets.
The empirical performance of the normalizing flow based variant of our method is displayed in Figure 11 As can be seen, the distributional Sobolev variants we propose still outperform the deterministie variant proposed in MAGE (D'Oro & Jaskowski 2020 under the stochastic settings introduced in Section 6.2
Mormel
Mattiplictiiee nose
Doussken nobs
400
000
10.50
dod
150 ate
280 ole
A 200
Onv mage
Dy mag
Er Pags
8000
10:00
120.000
17.00
27.00
25.00
tin
Munanolteva
#uber Sobalov (MAGE)
MDO Sobo
MSMDO Sobo
Figure 11: Full curves DSDPG and MAGE D Oro & Jaskowski 2020 using a RealNVP Dinh et al. 2016) normalizing flow differentiable world model.
59

---

Published as a conference paper at ICLR 2026
M.2 OVERESTIMATION BIAS
In our inplementation we rely on the TQC trick (Kuznetsov et al. 2020), which mitigates overesti- mation by discarding the largest values in the target distribution. This mechanism introduced and discussed in Section 1.1] is essential for stable learning: in our ablations, removing truncation leads substantially degraded performance. For completeness, we also remove the double-estimation correction from MAGE (D'Oro & Jaskowski 2020). These ablations show that controlling over- estimation is critical for reliable training. The corresponding results are reported Figures 12a and 12b. which present the final evaluation performance and the normalized Area Under the Curve with and without the overestimation biaiionoobbsssssaa countermeasure.
5903
2003
Womm
Mullil
Fik
Wullal salive
Geuailir robe
Wallerid-2
4303
Normal
Multet catiee nobl
Geomaa robs
3003
4003
2033
Manmoal
Miltgg cative
Dayodaa rete
Matrul
Ganoiin 19b6
Hitbe Saltolew IMAGE)
VID Satocley
VDMHO Subblley
2771 Forned undernntiraaiion bacc ITDATIOCI
No Forted underest mation Dan
E77 Forced undernntimationn biss (TDATTCC)
Iio forted underedd mation
(a)
(b)
Figure 12: Comparison of MAGE D'Oro & Jaskowski 2020) and DSDPG with and without over- estimation bias measures. We removed the double estimation of MAGE borrowed from (Fujimoto et al. 2018) and the TQC trick (Kuznetsov al. 2020) of our DSDPG method.
M.3 MULTIQUADRIC KERNEL BANDWIDTH
We recall the multiquadric kernel
1/ (x.y) v) = 121l v112.
As shown in Killingberg & Langseth 2023), selecting proper value for h is critical as it affects both the expressiveness of the kernel and the numerical stability of its MMD estimator: a too-small h yields nearly constant kernel, reducing its discrimination power, while a too-large h increases the gradient magnitude, scaling on the order of h, which can lead to exploding gradients and training oscillations unless mitigated by techniques such kernel rescaling gradient clipping.
In Section 6.2 the default value for was 100 which worked reliably. Here we test the values h. = 10 and h 250. The results are depicted in Figure 14 and 13 where we display the final evaluation performance and the evaluation curves. The comparison is made on the environments with and without multiplicative noise (as in Section 6.2). As can be overall, the performance of MMD and DSDPG (MMD Sobolev) seems to be quite insensitive to the kernel bandwidth. Our choice of h 100 seems to be robust.
M.4 NOISE SCALE
Since the noise scale we exposed in Section 6.2, which we called medium, was already sufficient to observe gap between DSDPG and the baseline, tried to increase the noise scale. We moved from n ~ (0.8, 1.2) to n ~ 4/007, 1.3). These changes already made the tasks harder. The results are the good performance of DSDPG (MMD Sobolev) depicted
displayed in Figure 15 As can be in Section 6.2 can be extended to the larger noise scale, as our method maintains a consistent gap against the baselines (especially MAGE).
60

---

Published as a conference paper at ICLR 2026
redken
recken
vodex
redoe
AdM
wedee
Figure 13: Full curves of DSDPG against baselines on 5 MuJoCo environments.
HaPCheetth-12
Wolter2dw2
beut3
300
PM0
300
PN0 9Y
MAO
MDC Sebee
Pio
400
Figure 14: Comparison of bandwidths under the multiquadric kernel for MMD distributional RI. and Sobolev MMD distributional RL. Bar plots of the final averaged evaluation reward. Mcan over 5 seeds.
61

---

Published as a conference paper at ICLR 2026
PAD
(TS8)
2000
70.0
3000
1906
Stepx
Steps
Humanoid-v2
no noise
medium noise
large noise
3000
2000
2009
1009
106009
101000
11600
2000
16000
200t
Wel
100009
Nomm
Walker2d-v2
large
2200
125000
1500
Strps:
Strge
Steps
Figure 15: Comparison of multiplicative noise scales. Averaged evaluation sum of rewards over seeds.
62

---

Published as a conference paper at ICLR 2026
M.5 NUMBER OF SAMPLES
Sensitivity to the number of samples is an important design question as the cost of generating sam- ples will scale linearly and the cost of estimating MMD will scale quadratically. It also interesting to verify how sensitive our method is to that parameter. The results displayed in Section 6.2 used 10 samples
model the Sobolev distributions. Here we additionally test with and 25 samples. In Figure 17 and T6 we display the final evaluation returns and evaluation curves while varying the number of samples.
can be seen, there no clear trend to be found. The Ant-v2 environment seems to be particularly sensitive to this parameter. Apart from this example we observe our method robust to the number of samples as at least two different values maintained performance.
Marcheetahvg
Walle
Hoppare22
Huranolde2
Figure 16: Comparison of number of samples used in the MMD methods (usual MMD and DSDPG denoted as MMD Sobolev). Evaluation curves are mean over 5 seeds.
M.6 WORLD MODEL CAPACITY
The world model size is, along the number of samples, one of the primary parameter driving the overall cost of our method. As for MAGE D'Oro & Jaskowski (2020). we backpropagate through the world model to infer the gradients of the target. Here change the width of each neural network in the cVEE world model from 1024 to 256. We observe in Figure 18 on the single Walker2d-v2 environment we tested, rather than being insensitive, it seems the choice of 1024 was sub-optimal for this enviromment, This suggests our method is robust to this design choice. This should be evaluated on more demanding (higher dimensional) environments (Ant-v2, Humanoid-v2).
N LLM USAGE
We used an LL.M-based assistant to support the preparation of this paper. In particular, it was employed to (i) rephrase draft paragraphs for clarity and suggest alternative framings of related work, (ii) format proofs, explore directions, and verify intermediate steps, (iii) assist in debugging code, (iv) suggest LaTeX equation formatting, and (v) help identify relevant theoretical results in preceding works. All core research contributions, including the development theoretical results, algorithms, and experiments, were carried out by the authors.
63

---

Published as a conference paper at ICLR 2026
Woiber2l-v2 +2
L4 tate
Mooper-e2
Ant-22
Humanald--22
Figure 17: Comparison of number of samples used in the MMD methods (usual MMD and DSDPG denoted as MMD Sobolev). Final evaluation performance as mean over 5 seeds.
Walker2d-v2 medium noise
MMD, large
MMD, small
Average evaluation reward
4000
MMD Sobolev, large
MMD Sobolev, small
3000
2000
1000
0
25000
50000
75000
100000
125000
150000
Steps
Figure 18: Comparison of world model size for MMD distributional RL and Sobolev MMD distri- butional RL. Bar plots of the final averaged evaluation reward. Mean over 5 seeds.
64


Reviews:
[{"review_id": "aSvFmUVMX4", "reviewer": "Reviewer_Hcbc", "summary": "This paper addresses a limitation of existing gradient-aware reinforcement learning methods, which improve sample efficiency by incorporating value gradients into critic training. The authors identify that current approaches treat these gradients as deterministic quantities. Thus, they extend distributional RL to model the joint distribution of both the cumulative return and its action-gradient. The proposed algorithm, Distributional Sobolev Deterministic Policy Gradient is built upon three key technical components. First is a novel Sobolev Bellman operator, which forms the theoretical core by bootstrapping the joint distribution of returns and their gradients. Second, to make this operator practical in non-differentiable environments, the method employs a conditional Variational Autoencoder enabling gradient backpropagation through the learned dynamics. Third, the critic is trained to minimize the discrepancy between its predicted distribution and the bootstrapped target distribution.", "questions": "What are the reasons for choosing the incomplete operator for DSDPG?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Extending the idea of distributional RL to action-gradient is interesting. By proposing to learn the distribution of the gradient, the work offers a principled and more complete model of credit assignment under uncertainty. This moves the field beyond the limitations of prior methods like MAGE, which are confined to modeling the gradient's expectation and are thus vulnerable to noise.\n- The formal derivation of the Sobolev Bellman operator is solid.\n- Except for theoretical contributions, the paper also provides comprehensive experimental results.", "weaknesses": "- As the authors stated, one weakness of DSDPG is its computational complexity.\n- There is a disconnect between the theoretical motivation for the MSMMD metric and its empirical utility. The paper introduces MSMMD primarily to obtain a provable contraction for the Sobolev Bellman operator under a tractable, sample-based metric. Standard MMD does not offer this same general guarantee. However, the empirical results across both the toy problem and the MuJoCo benchmarks show that the MSMMD variant provides at best a marginal performance improvement over the simpler MMD variant, and in some cases performs identically.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EXduBwE0qJ", "reviewer": "Reviewer_Dy2h", "summary": "This paper considers distributional reinforcement learning, and describes how to incorporate gradient information into distributional temporal-difference learning. Not only is it theoretically well-motivated, it also describes practical means of implementation even in challenging, realistic settings such as non-differentiable environments.", "questions": "No further questions.", "limitations": "", "rating": 10, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- This paper makes foundational contributions to the important and timely topic of distributional RL.\n- It is exceptionally well-written, with both pedagogical clarity and mathematical precision (admittedly, it gets a bit dense at times, but I think that is unavoidable). \n- The positioning against the Related Work is detailed, clear, and well articulated. Furthermore, it gives a balanced and honest discussion of its limitations.\n- I appreciate the pedagogical toy problem\n- The extensive appendices provide detailed support (mostly proofs) for the claims in the main paper.", "weaknesses": "Honestly, I can't think of any major weaknesses. Of course, it would have been nice to see clear improvements over the baselines on an established benchmark. Still, I think the theoretical and algorithmic contributions by far outweigh the importance of such experimental results. \n\nMinor things: \n- Some of the figures (and especially the text therein) are tiny and require zooming.\n- As a matter of personal preference, I think it's neater to emphasize text with italics rather than bold font, but I'll leave that for your discretion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TtbmY86Bvs", "reviewer": "Reviewer_dVoA", "summary": "This paper seeks to improve the sample efficiency and performance of actor critic RL algorithms in stochastic environments by estimating the distribution of  action gradients through the critic. To do so it introduces the distributional modeling of returns and their gradients, and also details an approach to use it for actor critic training, which it terms distributional Sobolev training. The practical instantiation uses a form of maximum mean discrepancy to estimate the action gradients via a novel Bellman operator, and uses a variational autoencoder to estimate the one-step transition function.\n\nEvaluations are done on a toy domain to verify that the approach works as intended and then in MuJoCo environments with added stochasticity to evaluate the approach on a standard benchmark.", "questions": "* Are practical considerations of implementing the idea potentially holding back the empirical results? Could the proposed approach be used for other robotics tasks like OGBench with added stochasticity?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* The proposed idea is intellectually stimulating and pushes the frontier on how actor critic algorithms work. It presents a novel idea and sound theoretical analysis to justify it.\n* The proposed practical algorithm takes reasonable approaches for the required theoretical estimates, such as MMD and MSMMD.\n* The use of the toy domain is useful to study if the proposed technique can deal with increased stochasticity as claimed\n* The statistical bounds used in the evaluations are heartening and appreciated.\n* The theoretical analysis seems to be sound\n* The idea would be of interest to researchers who focus on dealing with stochasticity in the environment.", "weaknesses": "* The use of the VAE is a practical requirement to deal with environments without differentiable dynamics. Using a toy environment where the gradients of the environment are known would have been even better to evaluate how well the idea would work if the dynamics did not have to be estimated.\n* While the results in the toy domain are impressive, it is unclear why they do not translate as well to the Mujoco environments, where MAGE remains within the margin of error. This difference is not a dealbreaker, since the trends across environments shows that the proposed techniques outperform MAGE across multiple environments.\n* Nit: Ordering the methods in Figure 3 so the proposed methods are the rightmost ones and the ones on the left are baselines would make them much easier to read", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dJus7ty5TK", "reviewer": "Reviewer_Z5td", "summary": "This paper proposes a new deterministic policy gradient method in the continuous action space by using Sobolev training, which additionally incorporates the gradient information in the distributional learning. Theoretically, the authors show the contractivity with Max-sliced MMD and Wasserstein distance under some assumptions. The dynamics and reward transitions are also needed to learn to derive a tractable algorithm. Experiments are conducted in a toy example and several Mujoco environments.", "questions": "1.What is the environmental non-differential regarding?\n\n2.The formulation in Eq.10 is questionable. Note that a is a constant value as the initial action, why do you mean the gradient is taken regarding a constant? In addition, the notation of cumulative rewards does not involve the action. Since it is a continuous action space, what does $|A|$ mean? Essentially, it should be $\\mathbb{R}^d$.\n\n3.What is the computational cost when we use max-sliced MMD? It seems that it is very high, but we are disappointed to find that the improvement is insignificant even with the higher computational cost.\n\n4.Why do we need to incorporate the TQC trick? Is there any ablation study in experiments compared with other baselines?\n\n5.The introduction of CVAE is less convincing. Why do we not directly train a diffusion model?\n\n6.The theory and the practical algorithm seem to have a gap. The finally proposed algorithm does not enjoy a theoretical contraction guarantee.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "1.Incorporating the action-gradient in the distributional learning seems novel.\n\n2.Some theoretical results are provided and also many details have been discussed, which are deferred to Appendix.", "weaknesses": "1. **The motivation is not clear**. Why is introducing action gradients beneficial in distributional learning? Directly combining Sobolev training seems straightforward. A more profound analysis is needed to strengthen the motivation.\n\n2. **Many theoretical results are questionable with impractical or unclear conditions.** I push back against the statement in line 142 he Lipschitz continuity of $\\pi$ typically holds when using neural-network function approximation, which is not true. The Lipschitz constant of deep neural networks is often unboundedly large, which depends on the spectral norms of the weight matrix that is often unbounded during the training. The contraction in the theorem only holds when $\\gamma \\kappa <1$, which is hard to justify for any practical environment. This renders the contraction vacuous and meaningless. What are the mild conditions in Theorem 2, which are very unclear. What is $\\kappa$ in a rigorous way?\n\n3. **Limited experiments and insignificant improvements**. In distributional RL literature, published papers have done extensive experiments with significant improvement, making the proposed algorithm deployable. However, the improvement in a toy example is too weak, and the improvement in the Mujoco environments is very limited. Although the robustness setting is helpful to some extent, it is unsatisfactory to find that the proposed algorithm only matches the performance with the baselines, even though it has additionally learned the environment and there is a much higher computational cost.\n\n4. **The writing is not clear.** Many sentences are in bold, but they are hard to follow. For instance, what do the authors want to express in Line 203? It is in bold but may not be strongly related to the other parts of the paper. A lot of details are provided in the appendix, which are very distracting. The main content of the paper is not even self-contained. For example, it is surprised that many key details in Theorem 2 are deferred to the appendix, which largely undermines the readability of this paper. I could not understand why we should not expect the proposed algorithm to be superior over all the environments in general.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)