Global Convergence of Wasserstein Policy Gradient for Entropy-Regularized Reinforcement Learning

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Global Convergence of Wasserstein Policy Gradient for Entropy-Regularized Reinforcement Learning

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Global Convergence of Wasserstein Policy Gradient for Entropy-Regularized Reinforcement Learning Zhaoyu Zhu

[email protected]

Shanghai Jiao Tong University

Rui Gao

[email protected]

arXiv:2605.26078v1 [cs.LG] 25 May 2026

The University of Texas at Austin

Shuang Li

[email protected]

The Chinese University of Hong Kong, Shenzhen

Abstract Wasserstein policy gradient (WPG) is a policy optimization method for reinforcement learning (RL) that exploits the optimal-transport geometry of action distributions. For the entropy-regularized RL objective, WPG evolves each state-conditional policy by transporting it along the action gradient of the soft Q-function together with a Langevin-type diffusion. Despite its appeal for continuous-control problems, its global convergence properties remain poorly understood. Standard Langevin analyses do not directly apply, because the RL objective depends on the policy through the Bellman recursion rather than through a static convex functional, and the Langevin drift is determined by the soft Q-function, whose regularity must be controlled along the policy iterates. In this paper, we develop a global convergence theory for WPG by exploiting the Bellman structure of entropy-regularized RL. We show that the role usually played by convexity can be replaced by a Bellman-based argument: the soft Bellman residual admits a statewise KL representation with respect to a Gibbs policy; Bellman contraction relates this residual to the global optimality gap; and a Bellman resolvent identity connects value improvement to relative Fisher information. Combined with a uniform log-Sobolev inequality (LSI) for the evolving Gibbs family, these ingredients yield a distributional Polyak–Lojasiewicz condition. We further establish the regularity and uniform bounds needed to control the discretization error, thereby obtaining geometric contraction up to a discretization bias. Conceptually, our analysis shows that although entropy-regularized RL is not convex in the usual flat sense, the Bellman recursion induces a favorable Polyak–Lojasiewicz-type (PL) geometry that supports global convergence of WPG.

1 Introduction Policy gradient (PG) methods are among the most widely used algorithms in reinforcement learning (RL). Beginning with REINFORCE (Williams, 1992) and the policy gradient theorem (Sutton et al., 2000), classical PG performs Euclidean gradient ascent over a parameterized policy class. While simple and scalable, Euclidean PG is sensitive to step sizes and parameterization, which has motivated geometry-aware alternatives. Natural policy gradient replaces the Euclidean metric with the information geometry induced by the Fisher information (Amari, 1998; Kakade, 2001), and is closely related to trust-region (TR) methods. In deep RL, TRPO constrains the average KL divergence between successive poli1

cies (Schulman et al., 2015), while PPO provides a practical surrogate through clipping or adaptive KL penalties (Schulman et al., 2017). These developments underscore the central role of geometry in policy optimization. Beyond information geometry, optimal transport (OT) provides a different geometry for policy space, one that is particularly natural for distributions over continuous action spaces. Existing Wasserstein policy optimization methods can be viewed as different implementations or approximations of Wasserstein gradient flow: through JKO schemes and particle approximations (Zhang et al., 2018), kernelized Wasserstein natural-gradient approximations (Arbel et al., 2020; Moskovitz et al., 2020), trust region formulations (Terpin et al., 2022; Song et al., 2024), and Wasserstein gradient updates projected onto explicit (Pfau et al., 2025) or implicit (Zhu et al., 2026) policy classes. These methods share the same underlying mechanism: policy improvement is driven by transport in action space. Entropy regularization is another powerful scheme in modern continuous-control RL. Maximum-entropy RL augments reward maximization with an entropy term on the policy distribution over actions, promoting exploration and robustness and leading to soft Bellman operators and Gibbs-type policy improvements (Ziebart, 2010; Haarnoja et al., 2018; Geist et al., 2019). From the OT perspective, this entropy term naturally induces diffusion. Consequently, Wasserstein policy optimization for the entropy-regularized objective takes the form of a drift-diffusion, or Fokker–Planck, equation, whose particle representation is a Langevin-type actor update. More concretely, the Wasserstein Policy Gradient (WPG) flow reads ∂t πt (a | s) = −diva πt (a | s)∇a Qπt (s, a) − τ log πt (a | s) , and its discrete-time Langevin counterpart is p Ask+1 = Ask + η∇a Qπk (s, Ask ) + 2τ η ξk+1 ,

ξk+1 ∼ N (0, Id ),

(1)

where πk (· | s) = Law(Ask ). Thus, each state-conditional policy is transported along the action gradient of the current soft Q-function, while entropy regularization contributes to the Langevin diffusion. Despite the algorithmic appeal of WPG, its global convergence behavior remains largely open. Existing global analyses of PG or natural PG typically rely on Euclidean or KL/Bregman geometry in policy space, using mirror-descent identity to relate policy updates to global suboptimality (Agarwal et al., 2021; Bhandari and Russo, 2024; Cen et al., 2022; Bhandari and Russo, 2021; Khodadadian et al., 2021, 2022; Yuan et al., 2023; Xiao, 2022; Lan, 2023; Zhan et al., 2023; Lan et al., 2023). These arguments do not transfer to Wasserstein geometry, which is non-Bregman and lacks the pointwise three-point structure underlying KL-based analyses. Existing OT-based convergence results also do not directly help for the Langevin WPG update in (1). The convergence result in Zhang et al. (2018) is asymptotic and places the entropy regularization on the policy parameters rather than on state-conditional policy distributions. The trust-region analysis of Song et al. (2024) applies to finite tabular MDPs without entropy regularization, and Zhu et al. (2026) studies a JKO-type implicit proximal scheme. None of these analyses transfers to Langevin dynamics. Our analysis is closely related to the broader literature on mean-field Langevin dynamics, which has been used to study training dynamics of large interacting particle systems and over-parameterized neural networks (Chizat and Bach, 2018; Mei et al., 2018; Nitanda 2

et al., 2022; Chizat, 2022; Cai et al., 2025). However, global convergence arguments in that literature typically rely on a convex objective and convexity-analysis based argument. RL objective has a different structure: the policy enters the objective through the Bellman recursion which destroys the convexity, and the soft Q-function entangles the reward and entropy term which make the regularity of the Langevin drift not taken for granted. Contribution. Our main contribution is a non-asymptotic global convergence analysis of WPG for entropy-regularized continuous-action RL. Rather than relying on the convexanalysis machinery commonly used in standard Mean-field Langevin dynamics, we exploit the Bellman structure of the RL objective. The analysis is built on the following ingredients. • Value improvement-to-Fisher information via a Bellman resolvent identity. We show that policy-space dissipation propagates through the Bellman equation via a resolvent identity. This allows us to lower bound value improvement by Fisher-information dissipation, replacing the standard energy-dissipation identity used in convex settings. • Optimality gap-to-KL via Bellman residuals and contraction. We represent the statewise soft Bellman residual as a KL divergence between the current policy and its associated Gibbs policy. Bellman contraction then relates this residual to the global optimality gap in sup-norm, replacing the entropy-sandwich arguments used in convex settings. • Uniform control of the moving Gibbs family. To exploit the commonly used uniform logarithmic Sobolev inequalities to connect the KL divergence (controlling the optimality gap) with the Fisher information (controlling value improvement), we establish uniform bound estimates along the WPG iterates. We first illustrate these ideas in the continuous-time WPG flow, where the Bellman residual, contraction, resolvent mechanisms appear most transparently. Then we provide careful analysis for the discrete-time Langevin update (1), which requires additional regularity estimates for the Wasserstein gradient and a one-step interpolation argument (Vempala and Wibisono, 2019) to control time-discretization error. To this end, we establish refined uniform a priori bounds on value functions, soft Q-functions, drift Lipschitz constants, action moments, drift moments, and the relevant KL quantities. Combining these ingredients yields geometric contraction of the global optimality gap up to a discretization bias. Organization. Section 2 introduces the entropy-regularized discounted RL model, the soft Bellman operators, and the Wasserstein/Langevin policy update. Section 3 presents the continuous-time analysis, isolating the core Bellman mechanisms underlying the proof. Section 4 proves the finite-time convergence theorem for the discrete-time update, with emphasis on the additional regularity estimates and discretization-error control. The appendices contain auxiliary results and detailed proofs.

2 Model Setup We consider an infinite-horizon discounted Markov decision process M = (S, A, P, r, γ, ρ0 ), 3

where S is the state space, A is the action space, γ ∈ (0, 1) is the discount factor, ρ0 ∈ P(S) is the initial state distribution, P (ds′ | s, a) is a Markov transition kernel on S, and r : S × A → R is a measurable reward function. A (stationary Markov) policy is a conditional distribution π(da | s) that admits a density, still denoted π(a | s), with respect to Lebesgue measure on A. Given π, a trajectory (sn , an )n≥0 is generated by s0 ∼ ρ0 , an ∼ π(· | sn ), and sn+1 ∼ P (· | sn , an ). We write Eρ0 ,π [·] for expectations under this rollout distribution. The normalized discounted state occupancy of π is the probability measure dπ ∈ P(S) defined by X γ n Pρ0 ,π (sn ∈ B), B ∈ B(S). dπ (B) = (1 − γ) n≥0

We assume throughout that the initial distribution ρ0 has full support on S. Hence, for every policy π, X dπ (B) = (1 − γ) γ t Pρ0 ,π (st ∈ B) ≥ (1 − γ)ρ0 (B), B ∈ B(S). t≥0

Thus dπ has full support whenever ρ0 does. This ensures that the statewise policy-gradient updates below are specified on the same state space on which the sup-norm value bounds are proved. We use the following standing measurability convention. The model kernels and their action derivatives admit jointly measurable versions such that all displayed integrals are jointly measurable. Under this convention, Bellman fixed points are selected as bounded measurable fixed points of contraction maps, and Gibbs and WPG kernels are taken in their induced jointly measurable versions. We work with an action space A = Rd . We consider the regularized reward r̃(s, a) = r(s, a) −

β ∥a∥2 , 2

where β > 0 is the quadratic action penalty. This penalty term stabilizes the training (Lillicrap et al., 2015; Brockman et al., 2016; Tunyasuvunakool et al., 2020), and is also exploited in standard mean-field Langevin dynamics analysis (Yamamoto et al., 2024), closely related to the dissipativity condition that helps to control the KL divergence along the trajectory. Fix a temperature τ > 0, we study the discounted regularized objective   X J(π) = Eρ0 ,π  γ n (r̃(sn , an ) − τ log π(an | sn )) , n≥0

restricting attention to policies for which the above expectation is finite. The quadratic action penalty is essential on the unbounded action space. Without this term, the entropy bonus can be made arbitrarily large by spreading the policy mass over larger regions of Rd , so the entropy-regularized policy-improvement objective may be unbounded above. The quadratic penalty rules out this degeneracy. Equivalently, the quadratic term induces the Gaussian reference β 2πτ d/2 −1 2 . ρβ (a) = Zβ exp − ∥a∥ , Zβ = 2τ β 4

For the discrete-time result, we work on the admissible class n o Π := π : π(· | s) ≪ da, (s, a) 7→ π(a | s) jointly measurable, sup KL(π(· | s)∥ρβ ) < ∞ . s∈S

Finite relative entropy gives finite second moments; see Lemma E.1 in Appendix E. 2.1 Soft value functions and Bellman operators Given a policy π, define the soft value function   X γ n r̃(sn , an ) − τ log π(an | sn ) | s0 = s , V π (s) = E  n≥0

so that J(π) = value function

π (s)ρ (ds) whenever the integral is well-defined. Define the state-action 0

Q (s, a) = r̃(s, a) + γ

V π (s′ )P (ds′ | s, a).

Then V π satisfies the soft Bellman identity Z π V (s) = (Qπ (s, a) − τ log π(a | s)) π(a | s)da. A

For a fixed policy π, define the policy evaluation operator acting on bounded measurable V : S → R by Z Z π ′ ′ (T V )(s) = r̃(s, a) + γ V (s )P (ds | s, a) − τ log π(a | s) π(a | s)da. (2) A

Under standard measurability and integrability conditions, T π is a γ-contraction in ∥ · ∥∞ and its unique fixed point is V π = T π V π . For a bounded candidate value V , write Z QV (s, a) := r̃(s, a) + γ V (s′ )P (ds′ | s, a). S

Define the soft optimality operator by Z ⋆ (T V )(s) = sup µ

QV (s, a) − τ log µ(a) µ(a)da,

where the supremum ranges over action densities for which the expression is finite. The supremum is attained uniquely at a Gibbs density, and therefore Z QV (s, a) exp(QV (s, a)/τ ) ⋆ (T V )(s) = τ log exp da, G[V ](a | s) = R . τ Rd Rd exp(QV (s, ã)/τ )dã Moreover, T ⋆ is a γ-contraction in ∥ · ∥∞ and has a unique bounded fixed point V ⋆ = T ⋆ V ⋆ . We set Q⋆ := QV ⋆ and π ⋆ := G[V ⋆ ]. For each (π, s), define the Gibbs density by pπs (a) := G[V π ](a | s) = R

exp(Qπ (s, a)/τ ) , π Rd exp(Q (s, ã)/τ )dã

a ∈ Rd .

(3)

The following identity relates the Bellman residual to the KL divergence between the current policy and the Gibbs density. 5

Lemma 2.1 (Bellman residual identity) For every s ∈ S, (T ⋆ V π )(s) − V π (s) = τ KL(π(· | s)∥pπs ) ≥ 0, where pπs is the Gibbs density defined in (3). 2.2 Wasserstein Policy Gradient Fix a policy π. The first-variation calculation in Appendix A.2 shows that, for statewise transport perturbations of the action law, the ascent direction is ∇a Qπ (s, a) − τ log π(a | s) . The discounted occupancy dπ enters this calculation through the usual policy-gradient weighting over states, as in the policy-gradient theorem and trust-region/natural-gradient policy methods (Sutton et al., 2000; Kakade, 2001; Schulman et al., 2015). Taking the Wasserstein gradient of this first variation state by state gives the Wasserstein policygradient flow ∂t πt (a | s) = −∇a · πt (a | s)∇a Qπt (s, a) − τ log πt (a | s) .

(WPGF)

Using the Gibbs density pπs t in (3), this can be rewritten as πt (a | s) ∂t πt (a | s) = τ ∇a · πt (a | s)∇a log πt . ps (a) From an equivalent particle viewpoint, for each fixed state s, the above Fokker–Planck equation can be viewed formally as the forward equation of the nonlinear Langevin diffusion dAst = ∇a Qπt (s, Ast ) dt +

√

2τ dBts ,

πt (· | s) := Law(Ast ).

Indeed, since τ ∇a log pπs t (a) = ∇a Qπt (s, a), Applying an explicit Euler–Maruyama step to this statewise Langevin diffusion gives the discrete-time update below. The discrete-time counterpart of (WPGF) is the Langevin policy update Ask+1 = Ask + η∇a Qπk (s, Ask ) +

p 2τ η ξk+1 ,

ξk+1 ∼ N (0, Id ),

(WPGD)

where Ask ∼ πk (· | s) and πk+1 (· | s) := Law(Ask+1 ). Thus WPGD is a statewise policy-space update specified for every s ∈ S. Under the fullsupport convention above, the discounted occupancy weighting does not remove any state from the gradient calculation, and the convergence theorem below controls the resulting value gap in sup-norm over the same state space. 6

3 Continuous-Time Analysis: Bellman Geometry This section presents a formal continuous-time analysis of the Wasserstein policy gradient flow (WPGF). Its purpose is to make transparent the Bellman mechanisms that replace the standard mean-field Langevin dynamics proof chain. We do not claim here a complete wellposedness theorem for the nonlinear Fokker–Planck equation. Instead, the identities below are stated under a smooth-flow regularity convention that supplies the differentiability, integrability, and boundary-decay properties needed for the calculation. Why the standard mean-field Langevin dynamics proof chain does not directly apply. For a mean-field Langevin dynamics problem, exponential convergence is often proved through the schematic chain (Nitanda et al., 2022; Chizat, 2022): (b) (c) d (a) F(µt ) = − I(µt ∥νt ) ≲ − KL(µt ∥νt ) ≲ − F(µt ) − F ⋆ , dt

(4)

where (a) is the gradient-flow dissipation identity, (b) follows from a uniform LSI, and (c) is through a convex-analysis based comparison from local KL to global suboptimality. In entropy-regularized RL, these relationships do not hold in these forms. First, (a) is altered by the Bellman recursion: differentiating V πt also differentiates Qπt through the value function, so the value derivative is not merely a statewise Fisher-information dissipation. Lemma 3.1 replaces this step by the value-resolvent identity (I − γP πt )V̇ πt = gt , with gt proportional to a statewise Fisher-information term. Second, (b) is not directly available from a convex functional, because the RL objective is nonconvex due to the Bellman recursion of the state-dependent policy distribution. Lemma 3.2 and 3.3 will derive a uniform LSI from uniform value/Q bounds. Third, (c) cannot be obtained from convexity due to the nonconvexity and aggregation over states in the RL objective. Instead, Lemma 2.1 identifies τ KL(πt (· | s)∥pπs t ) with the statewise Bellman residual, and Bellman contraction in the proof of Theorem 3.1 will compare this residual at the worst state directly with the ∥ · ∥∞ -value gap. Smooth-flow regularity convention. Throughout this section only, the flow (WPGF) is considered for smooth positive solutions πt (· | s) satisfying the following properties on every finite time interval. For each state s, t 7→ V πt (s) is locally absolutely continuous, the density πt (· | s) has the moments, entropy, Fisher information, and decay at infinity needed for all displayed integrations by parts, and a 7→ Qπt (s, a) is sufficiently smooth. Differentiation under the action/state integrals is assumed valid. These assumptions are used only for the formal PDE calculation in this section for illustration and will not be used in the discrete-time analysis in Section 4. Assumption 3.1 (Boundedness and initialization) We assume the following throughout: (i) |r(s, a)| ≤ Rmax for all (s, a) ∈ S × A. β (ii) sups∈S KL(π0 (· | s)∥ρβ ) ≤ K0 , where ρβ (a) ∝ exp − 2τ ∥a∥2 . Our proof proceeds in three steps. First, we establish a resolvent identity for the value derivative along the flow (Lemma 3.1), which reveals a dissipation given by Fisher in7

formation relative to the moving Gibbs family pπs t and yields monotone value improvement. Second, we use this monotonicity together with bounded rewards and the initialization condition to obtain uniform-in-time bounds on V πt and Qπt (Lemma 3.2); these bounds imply that each pπs t is a bounded perturbation of a Gaussian, enabling a uniform log-Sobolev inequality along the optimization trajectory via Holley-Stroock Perturbation Lemma (Lemma 3.3). Third, we combine the uniform LSI with the Bellman-residual identity, Lemma 2.1, and Bellman contraction to convert KL dissipation into an exponentialdecay estimate for the optimality gap (Theorem 3.1). Note that RL policy differs from the standard mean-field Langevin dynamics in that it is concerned with a family of state-dependent policy distributions {π(· | s)}s∈S rather than a single distribution. Thus, adapting the usual proof chain (4) to RL settings requires choosing how to aggregate statewise quantities. Another possible route is to aggregate the KL terms in (c) by a discounted state-visitation distribution via the performance difference lemma A.1, as in occupancy-weighted policy-gradient analyses. However, this creates a mismatch between the state distribution used to measure the optimality gap and the state aggregation naturally induced by the value derivative or resolvent identity in (a). Closing this gap typically requires density-ratio or coverage assumptions between state distributions (Xiao, 2022; Agarwal et al., 2021; Zhang et al., 2020; Ding et al., 2022). Our argument above circumvents this route by keeping the Bellman residual identity statewise and using Bellman contraction to lower bound the residual in terms of the sup-norm value gap, and thus does not require additional coverage assumptions. For a policy π, define the induced state kernel Z π ′ P (ds | s) := P (ds′ | s, a)π(a | s)da. Rd

Definition 3.1 (KL divergence and relative Fisher information) Let µ and ν be probability measures on Rd with µ ≪ ν. The Kullback–Leibler divergence of µ from ν is defined by Z dµ KL(µ∥ν) := log dµ. dν Rd The relative Fisher information of µ with respect to ν is defined by Z ∇ log

I(µ∥ν) := Rd

dµ 2 dµ. dν

Definition 3.2 (Log-Sobolev inequality) Let ν be a probability measure on Rd with positive density. We say that ν satisfies a log-Sobolev inequality with constant α > 0, or α-LSI, if I(µ∥ν) ≥ 2α KL(µ∥ν), ∀µ ≪ ν for every probability measure µ for which the two sides are well-defined. We first establish an identity for the time derivative of V πt along the flow (WPGF).

Lemma 3.1 (Monotonicity of V πt ) Under Assumption 3.1 and the smooth-flow regularity convention, for every t ≥ 0, (I − γP πt )

d πt V = gt , dt

gt (s) := τ 2 I πt (· | s)∥pπs t ≥ 0,

(5)

where I πt (· | s)∥pπs t is the relative Fisher information. Moreover, d πt V (s) ≥ gt (s) ≥ 0 dt

∀s.

(6)

Lemma 3.1 shows that the value derivative satisfies d πt V = (I − γP πt )−1 gt dt and is therefore pointwise nonnegative along the flow (WPGF). Note that this identity can be viewed as the statewise resolvent form of the entropy-regularized performance-difference lemma: integrating it against an initial distribution recovers the usual occupancy-weighted performance-difference formula, whereas keeping it pointwise is what enables the subsequent sup-norm contraction argument in (9) below without introducing state-distribution mismatch coefficients. Based on Lemma 3.1, we derive uniform-in-time bounds on the value function, which further imply uniform bounds on the action-value Q-function. Lemma 3.2 (Uniform bounds on V πt and Qπt ) Under Assumption 3.1 and the smoothflow regularity convention above, for all t ≥ 0, sup V πt (s) ≤ s∈S

Rmax + τ log Zβ . 1−γ

Moreover, inf V πt (s) ≥ inf V π0 (s) ≥

s∈S

−Rmax + τ log Zβ − τ K0 . 1−γ

Consequently, for all t ≥ 0, s ∈ S, and a ∈ A, −Rmax −

β β ∥a∥2 + γ inf V π0 (x) ≤ Qπt (s, a) ≤ Rmax − ∥a∥2 + γ sup V πt (x). x∈S 2 2 x∈S

(7)

Thanks to the bound (7) on the Q-function, the Gibbs densities pπs t are bounded perturbations of Gaussian densities, and thus satisfy log-Sobolev inequalities with uniform constants along the optimization trajectory due to the Holley-Stroock perturbation lemma. Lemma 3.3 (Trajectory-uniform LSI for Gibbs policies) Define Rmax + τ log Zβ Rmax − τ log Zβ + τ K0 Vmax = max , 1−γ 1−γ and Umax = Rmax + γVmax . 9

Then for every t ≥ 0 and s ∈ S, the Gibbs density pπs t satisfies the log-Sobolev inequality I(ν∥pπs t ) ≥ 2α KL(ν∥pπs t )

∀ν ≪ pπs t ,

with the uniform constant β 2Umax α = exp − . τ τ We now state the exponential-decay consequence of the preceding identities. Let V ⋆ be the unique fixed point of T ⋆ , and let π ⋆ denote the optimal Gibbs policy attaining V ⋆ . Since V ⋆ is the optimal value function and the flow remains in the admissible policy class, V ⋆ (s) ≥ V πt (s) for all s, t. Define the residual Rt (s) := (T ⋆ V πt )(s) − V πt (s) ≥ 0. By Lemma 2.1, Rt (s) = τ KL πt (· | s)∥pπs t .

(8)

Define Wt (s) := V ⋆ (s) − V πt (s) ≥ 0,

E(t) := ∥Wt ∥∞ .

Since V ⋆ = T ⋆ V ⋆ , T ⋆ is monotone, and T ⋆ is a γ-contraction in ∥ · ∥∞ , 0 ≤ V ⋆ (s) − (T ⋆ V πt )(s) ≤ γE(t)

∀s ∈ S, t ≥ 0.

(9)

Theorem 3.1 (Exponential convergence along the smooth WPG flow) Under Assumption 3.1 and the smooth-flow regularity convention above, for all t ≥ 0, J(π ⋆ ) − J(πt ) ≤ ∥V ⋆ − V πt ∥∞ ≤ e−2ατ (1−γ)t ∥V ⋆ − V π0 ∥∞ . Proof By Lemma 3.1, specifically (5) and (6), for almost every t, ∂t Wt (s) = −

d πt V (s) ≤ −gt (s) = −τ 2 I πt (· | s)∥pπs t . dt

By Lemma 3.3 and (8), 2α I πt (· | s)∥pπs t ≥ 2α KL πt (· | s)∥pπs t = Rt (s). τ Therefore, ∂t Wt (s) ≤ −2ατ Rt (s)

(10)

Next, using the definition of Wt and Rt , Wt (s) = V ⋆ (s) − (T ⋆ V πt )(s) + Rt (s). Hence, by (9), Rt (s) = Wt (s) − V ⋆ (s) − (T ⋆ V πt )(s) ≥ Wt (s) − γE(t). 10

(11)

Combining (10) and (11) gives the pointwise differential inequality ∂t Wt (s) ≤ −2ατ Wt (s) + 2ατ γE(t),

s ∈ S,

(12)

for almost every t. Multiplying (12) by the integrating factor e2ατ t and integrating from 0 to t yields Z t 2ατ t e Wt (s) − W0 (s) ≤ 2ατ γ e2ατ u E(u) du. 0

Equivalently, −2ατ t

Wt (s) ≤ e

Z t

e−2ατ (t−u) E(u) du.

(13)

e−2ατ (t−u) E(u) du.

(14)

W0 (s) + 2ατ γ 0

Taking the supremum over s ∈ S in (13) gives E(t) ≤ e−2ατ t E(0) + 2ατ γ

Z t 0

Define H(t) := e2ατ t E(t). Multiplying (14) by e2ατ t gives Z t H(t) ≤ E(0) + 2ατ γ

H(u) du. 0

By Gronwall’s inequality, H(t) ≤ E(0)e2ατ γt . Therefore, E(t) ≤ E(0)e−2ατ (1−γ)t = ∥V ⋆ − V π0 ∥∞ exp{−2ατ (1 − γ)t}. Finally, since ρ0 is a probability measure and Wt ≥ 0 pointwise, Z ⋆ J(π ) − J(πt ) = V ⋆ (s) − V πt (s) ρ0 (ds) ≤ ∥V ⋆ − V πt ∥∞ . S

This proves the theorem. Theorem 3.1 is a formal continuous-time convergence for smooth solutions of (WPGF). It identifies the rate obtained from the value-resolvent identity, uniform value/Q bounds, the uniform LSI, and the statewise Bellman-residual identity. Its role is explanatory: it isolates the Bellman-residual mechanism and motivates the discrete-time argument. After the first version of this paper was submitted for review, we became aware of a closely related note by Šiška and Zhang (2026) for finite-state continuous-action problems. Their analysis and resulting guarantee are structurally different from ours. In their costminimization notation, the argument uses an entropy-sandwich inequality to control the ⋆ scalar distributional gap V πt (ρ) − V π (ρ) and states a bound of the form Z π⋆ π⋆ π0 π⋆ πt V (s) − V (s) dρ (ds) e−2κατ (1−γ)t , V (ρ) − V (ρ) ≤ κ̄ S

where κ̄ := sups∈S dddρπ⋆ (s), κ := inf s∈S dddρπ⋆ (s) ≤ 1. Thus their estimate is distributionρ

dependent and requires coverage-type density-ratio control, in particular κ > 0, to obtain a positive exponential rate. By contrast, Theorem 3.1 proves a state-uniform bound and therefore controls J(π ⋆ ) − J(πt ) for any initial state distribution. The proof mechanisms are also different: rather than passing through an occupancy-weighted entropy sandwich, we use the exact statewise Bellman-residual identity and Bellman contraction to compare the residual directly with the ∥ · ∥∞ -optimality gap. This removes the density-ratio constants κ from the exponential convergence rate, although with a different initial gap. Finally, both their and our analyses involve formal calculations; the main rigorous guarantee in the present paper is the discrete-time Langevin analysis in Section 4.

4 Discrete-Time Global Convergence Analysis This section proves the non-asymptotic convergence guarantee for the explicit Langevin update (WPGD). The continuous-time section isolates the Bellman residual/resolvent mechanism. Apart from their counterparts in discrete time, an additional task is to control the finite-step Langevin discretization and the policy-dependent drift uniformly along the iterates. 4.1 Assumptions The assumptions below are the quantitative inputs for the finite-step analysis. To apply the one step drift interpolation argument, we need uniform control of the drift Lipschitz constants, action moments, drift second moments, finite KL quantities, and the LSI constants of the Gibbs family. Since Z bk (s, a) := ∇a Qπk (s, a) = ∇a r(s, a) − βa + γ V πk (s′ )∇a p(s′ | s, a)λ(ds′ ), S

these constants cannot be read off from the Wasserstein gradient of a static functional; they must be derived from the Bellman recursion and propagated along the WPG iterates. Assumptions 4.1-4.2, together with the a priori bounds that will be proved in Proposition 4.1, provide this uniform control. Assumption 4.1 (Bounded rewards and initialization) We assume: (i) |r(s, a)| ≤ Rmax for all (s, a) ∈ S × Rd . (ii) sups∈S KL(π0 (· | s)∥ρβ ) ≤ K0 . For constants used later, set M0 := sup E∥As0 ∥2 , s∈S

As0 ∼ π0 (· | s).

This quantity is finite by Lemma E.1 applied uniformly in s. In particular, applying that lemma with c = β/(4τ ) and using Assumption 4.1(ii) gives Z 4τ β 4τ d 2 M0 ≤ K0 + log exp ∥a∥ ρβ (a) da = K0 + log 2 . β 4τ β 2 Rd 12

Assumption 4.2 (Action regularity) There is a σ-finite reference measure λ on S such that P (ds′ | s, a) = p(s′ | s, a)λ(ds′ ). There are finite constants Lr , Lp , Gr , Gp such that: (i) For every s ∈ S, the map a 7→ r(s, a) is continuously differentiable and ∥∇a r(s, a)∥ ≤ Gr ,

∥∇a r(s, a) − ∇a r(s, ā)∥ ≤ Lr ∥a − ā∥.

(ii) For every s ∈ S and λ-a.e. s′ ∈ S, the map a 7→ p(s′ | s, a) is continuously differentiable, and Z ∥∇a p(s′ | s, a)∥λ(ds′ ) ≤ Gp ,

∥∇a p(s′ | s, a) − ∇a p(s′ | s, ā)∥λ(ds′ ) ≤ Lp ∥a − ā∥.

A concrete example satisfying the transition part of Assumption 4.2 is the following additive-noise model. Let S = Rm , let λ be Lebesgue measure, and suppose S ′ = F (s, a) + σζ, where σ ∈ Rm×m is invertible and ζ has a smooth density q. Write qσ (x) := | det σ|−1 q(σ −1 x) for the density of σζ, and assume Z

∥∇qσ ∥L1 < ∞,

∥D qσ ∥L1 :=

∥D2 qσ (x)∥op dx < ∞.

Assume further that, uniformly in s, ∥Da F (s, a)∥op ≤ M1 ,

∥Da F (s, a) − Da F (s, ā)∥op ≤ M2 ∥a − ā∥.

Then p(y | s, a) = qσ (y − F (s, a)) satisfies Assumption 4.2(ii). Indeed, ∇a p(y | s, a) = −Da F (s, a)⊤ ∇qσ (y − F (s, a)), so one may take Gp = M1 ∥∇qσ ∥L1 ,

Lp = M2 ∥∇qσ ∥L1 + M12 ∥D2 qσ ∥L1 .

Nondegenerate Gaussian noise is a special case, since its first derivative and Hessian are integrable. Together with the reward condition in Assumption 4.2(i), this gives a concrete class of controlled transition models covered by the theorem. Deterministic or degeneratenoise dynamics are not covered by this example. With these assumptions in place, write Wk := V ⋆ − V πk ,

Ek := ∥Wk ∥∞ ,

and define the current-policy Langevin drift bk (s, a) := ∇a Qπk (s, a). The current-policy Gibbs stationarity condition is τ ∇a log pπs k (a) = ∇a Qπk (s, a) = bk (s, a). 13

(15)

4.2 Proof sketch and main theorem 4.2.1 Bellman residual-to-optimality gap mechanism and uniform LSI We first record the discrete counterparts of the residual and resolvent mechanisms from Section 3. They show that the KL to the current-policy Gibbs law is exactly the Bellman residual, that this residual controls the worst-state value gap, and that decreasing the same KL produces value improvement after applying the Bellman resolvent. Lemma 4.1 (Bellman contraction of the residual) Fix k ≥ 0. Let Rk := (T ⋆ V πk ) − V πk . Then, for every s ∈ S, Rk (s) ≥ Wk (s) − γEk . Lemma 4.2 (Discrete Bellman resolvent) Fix k ≥ 0. Suppose πk and the WPG update πk+1 are admissible, with measurability covered by the convention in Section 2, and the KL terms below are finite. Then (16) (I − γP πk+1 ) V πk+1 − V πk = gk , gk (s) := (T πk+1 V πk )(s) − V πk (s), and, for every s ∈ S, gk (s) = τ KL(πk (· | s)∥pπs k ) − KL(πk+1 (· | s)∥pπs k ) .

(17)

As a consequence of these two lemmas, once pπs k satisfies an LSI with constant ᾱ, the Bellman residual yields the following PL-type residual-to-gap inequality at an ε-maximizer of Wk : I(πk (· | s)∥pπs k ) ≥ 2ᾱKL(πk (· | s)∥pπs k ) =

2ᾱ 2ᾱ Rk (s) ≥ (1 − γ)Ek − ε . τ τ

(18)

This is a PL inequality for the Bellman gap, not a consequence of flat convexity of J. The LSI used in (18) follows from the same Gaussian perturbation mechanism as in continuous time, but it must be made uniform along the discrete trajectory. The a priori bounds prove a uniform value bound; once ∥V πk ∥∞ ≤ Vmax , R r(s, a) + γ V πk (s′ )P (ds′ | s, a) eψk,s (a) ρβ (a) πk ps (a) = R ψ (ã) . , ψk,s (a) := τ e k,s ρβ (ã)dã The perturbation ψk,s may be nonconvex, but bounded rewards and the value bound give ∥ψk,s ∥∞ ≤ (Rmax +γVmax )/τ and hence osc(ψk,s ) ≤ 2(Rmax +γVmax )/τ . The Holley-Stroock bounded-perturbation principle (Holley and Stroock, 1987) then transfers the Gaussian LSI of ρβ to every pπs k uniformly. Lemma 4.3 (Uniform LSI) Fix k ≥ 0 and suppose ∥V πk ∥∞ ≤ Vmax . Then, for every s ∈ S, the current-policy Gibbs density pπs k satisfies I(µ∥pπs k ) ≥ 2α(Vmax )KL(µ∥pπs k ) where

∀µ ≪ pπs k ,

β 2(Rmax + γVmax ) . α(Vmax ) = exp − τ τ 14

4.2.2 Drift regularity and the discretization error We now control the finite-step error introduced by the Langevin discretization. As explained in Section 4.1, the drift Langevin interpolation requires uniform action-Lipschitz and dissipativity estimates for the policy-dependent drift. The following lemma converts the action-regularity assumptions and a value bound into those estimates. Lemma 4.4 (Bellman drift regularity) Assume Assumption 4.2. Fix k ≥ 0 and suppose ∥V πk ∥∞ ≤ Vmax < ∞. Then, for every s ∈ S, bk (s, a) = −βa + hk (s, a), with sup ∥hk (s, a)∥ ≤ G(Vmax ) := Gr + γGp Vmax .

(19)

s,a

Moreover, for all a, ā ∈ Rd , ∥bk (s, a) − bk (s, ā)∥ ≤ Lb (Vmax )∥a − ā∥,

Lb (Vmax ) := β + Lr + γLp Vmax .

The Lipschitz bound controls the Euler interpolation error; the dissipative decomposition controls the second moment and hence the squared drift moment. Together with the LSI from Lemma 4.3, they give the only place where discrete time loses exact monotonicity. The analytic part of the next lemma is based on a one-step fixed-drift interpolation. The full calculation is given in Appendix D.2. Lemma 4.5 (One-step KL contraction with discretization error) Fix k ≥ 0 and s ∈ S. Suppose pπs k > 0, log pπs k ∈ C 1 (Rd ), pπs k satisfies an LSI with constant α > 0, and the Gibbs-score identity τ ∇a log pπs k (a) = bk (s, a)

for all a ∈ Rd

holds. Suppose also that bk (s, ·) is Lb -Lipschitz and Ea∼πk (·|s) ∥a∥2 < ∞,

Ea∼πk (·|s) ∥bk (s, a)∥2 ≤ B 2 ,

Set δη (Lb , B 2 ) :=

KL(πk (· | s)∥pπs k ) < ∞.

L2b d 2 L2b B 2 3 η + η . 2 6τ

Then KL πk+1 (· | s)∥pπs k ≤ e−ατ η KL πk (· | s)∥pπs k + δη (Lb , B 2 ).

(20)

Lemma 4.6 (From KL contraction to Bellman improvement) Fix k ≥ 0 and s ∈ S. Suppose the admissibility and finite-KL hypotheses of Lemma 4.2 hold for (πk , πk+1 ). If, for some c ∈ [0, 1] and δ ≥ 0, KL(πk+1 (· | s)∥pπs k ) ≤ (1 − c)KL(πk (· | s)∥pπs k ) + δ,

(21)

then, with Rk (s) = τ KL(πk (· | s)∥pπs k ) and gk as in Lemma 4.2, gk (s) ≥ cRk (s) − τ δ.

(22)

In particular, Lemma 4.5 supplies (21) with c = 1 − e−ατ η and δ = δη (Lb , B 2 ) whenever its analytic hypotheses also hold. 15

With the uniform constants supplied by Proposition 4.1, the preceding lemmas yield a closed one-step recursion for the sup-norm gap. Throughout this argument, ᾱ denotes the trajectory-uniform LSI constant, L̄b the uniform Lipschitz constant of the drift, B̄ 2 the uniform second-moment bound for the drift, and δ̄η the corresponding one-step Langevin discretization error. Set cη := 1 − e−ᾱτ η . Then, for every s ∈ S, the estimates compose as follows: KL(πk+1 (· | s)∥pπs k ) ≤ (1 − cη )KL(πk (· | s)∥pπs k ) + δ̄η by Lemma 4.5, =⇒

gk (s) ≥ cη Rk (s) − τ δ̄η by Lemma 4.6,

=⇒

V πk+1 (s) − V πk (s) ≥ cη Rk (s) −

=⇒

by Lemma 4.2, Ek+1 ≤ 1 − (1 − γ)cη Ek +

τ δ̄η 1−γ

(23)

τ δ̄η 1−γ

by Lemma 4.1. Here the KL terms are statewise, Rk = T ⋆ V πk − V πk , and Ek = ∥V ⋆ − V πk ∥∞ . The substantive RL step is the middle of (23): a local KL decrease becomes a global valuegap contraction through the Bellman residual and the value resolvent, not through flat convexity of a probability functional. The implications in (23) reduce the convergence proof to estimates that hold uniformly along the WPGD trajectory. The constants entering this recursion are policy dependent, since both pπs k and bk are defined through Qπk and hence through the Bellman fixed point V πk . The next proposition supplies the required a priori estimates. Approximate value monotonicity controls the lower value envelope, the quadratic action penalty gives uniform action-moment bounds, and these estimates feed back into uniform drift regularity, finite KL controls, and a trajectory-uniform LSI for the moving Gibbs family. Proposition 4.1 (Uniform a priori bounds for the WPGD iterates) Under Assumptions 4.1 and 4.2, there exist explicit finite constants V̄ ,

L̄b ,

Ḡ,

M̄ ,

B̄,

ᾱ > 0,

η0 > 0,

depending only on the problem data and the initialization such that, for every 0 < η ≤ η0 , the WPG iterates are admissible and, for all k ≥ 0, ∥V πk ∥∞ ≤ V̄ ,

sup E∥Ask ∥2 ≤ M̄ , s∈S

sup Eπk ∥∇a Qπk (s, A)∥2 ≤ B̄ 2 . s∈S

Moreover, a 7→ ∇a Qπk (s, a) is L̄b -Lipschitz uniformly in (k, s), every current-policy Gibbs density pπs k satisfies LSI with constant at least ᾱ, and the generated policies satisfy sup sup KL(πk (· | s)∥ρβ ) ≤ K̄η ,

(24)

k≥0 s∈S

sup sup KL(πk (· | s)∥pπs k ) ≤ H̄η . k≥0 s∈S

(25)

All constants and the feasible step-size range are given in Appendix C. Proposition 4.1 provides the uniform constants needed to run the one-step recursion in (23) at every iteration. The value bound yields a common Holley–Stroock perturbation bound and hence a uniform LSI constant ᾱ; the drift Lipschitz and second-moment estimates give a uniform discretization error δ̄η = O(η 2 ); and the admissibility and finite-KL bounds ensure that the resolvent identity and the one step drift interpolation estimate apply throughout the trajectory. Substituting these uniform estimates into (23) gives a closed scalar recursion for the sup-norm optimality gap, which yields the geometric contraction up to the accumulated Langevin discretization bias stated in the main theorem. Theorem 4.1 (Global convergence of WPGD up to discretization bias) Under Assumptions 4.1 and 4.2, there exist explicit finite constants ᾱ, Cδ , η0 > 0, such that for every 0 < η ≤ η0 and every k ≥ 0, 2Cδ J(π ⋆ ) − J(πk ) ≤ ∥V ⋆ − V πk ∥∞ ≤ exp − 12 ᾱτ (1 − γ)ηk ∥V ⋆ − V π0 ∥∞ + η. ᾱ(1 − γ)2 Theorem 4.1 gives a finite-step geometric contraction for the WPGD update (WPGD), up to a fixed-step discretization bias. The drift KL estimate has an additive error δ̄η = O(η 2 ), and the scalar gap recursion divides this error by 1 − e−ᾱτ η ≍ ᾱτ η. The proof in Appendix D.4 keeps the sharper quantity δ̄η throughout and substitutes δ̄η ≤ Cδ η 2 only at the end. The theorem is closest in form to discrete-time mean-field Langevin guarantees, but the source of the contraction is different. Mean-field convex analyses convert KL decrease to objective decrease through flat convexity of a distributional objective (Nitanda et al., 2022; Chizat, 2022; Suzuki et al., 2023). Here the Gibbs law pπs k moves with the policy through the soft Bellman equation that depends on both the policy gradient and the KL regularization. The required PL-type relation is therefore reconstructed from the RL-specific ingredients. The result is also complementary to entropy-regularized natural policy gradient and policy mirror descent, where KL/Bregman geometry gives a closed-form policy-improvement step and Bellman monotonicity (Geist et al., 2019; Cen et al., 2022; Bhandari and Russo, 2021; Lan, 2023; Zhan et al., 2023). WPGD instead uses Wasserstein geometry: local progress is Fisher-information dissipation in action space, and the main technical burden is to show that this transport-based dissipation remains uniformly controlled under the policy-dependent Bellman recursion and relate it to the global optimality gap.

5 Conclusion We developed a convergence analysis framework for Wasserstein policy gradient in entropyregularized discounted RL with continuous actions. The proof combines Bellman-residual identities, value-resolvent relation, one-step interpolation, uniform value/Q-function and moment bounds, and uniform LSI for the iterated Gibbs family. These ingredients yield exponential convergence for the continuous-time WPG flow and geometric convergence for the discrete-time WPG Langevin dynamics up to a discretization bias. An important direction for future work is to integrate these policy-space guarantees with practical approximations, such as projection onto parametric policy classes and finite-sample critic estimation. 17

References Alekh Agarwal, Sham M. Kakade, Jason D. Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 22(98):1–76, 2021. Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10 (2):251–276, 1998. Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics ETH Zürich. Birkhäuser, 2 edition, 2008. Michael Arbel, Adil Salim, Anna Korba, and Zaid Harchaoui. Kernelized wasserstein natural gradient. In International Conference on Learning Representations (ICLR), 2020. Jalaj Bhandari and Daniel Russo. On the linear convergence of policy gradient methods for finite MDPs. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics (AISTATS), volume 130 of Proceedings of Machine Learning Research, pages 2386–2394. PMLR, 2021. Jalaj Bhandari and Daniel Russo. Global optimality guarantees for policy gradient methods. Operations Research, 72(5):1906–1927, 2024. doi: 10.1287/opre.2021.0014. Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540, 2016. Yang Cai, Siddharth Mitra, Xiuyuan Wang, and Andre Wibisono. On the convergence of min-max langevin dynamics and algorithm. In Proceedings of Thirty Eighth Conference on Learning Theory, volume 291 of Proceedings of Machine Learning Research, pages 677–754. PMLR, 2025. Shicong Cen, Chen Cheng, Yuxin Chen, Yuting Wei, and Yuejie Chi. Fast global convergence of natural policy gradient methods with entropy regularization. Operations Research, 70 (4):2563–2578, 2022. doi: 10.1287/opre.2021.2151. Lénaı̈c Chizat. Mean-field langevin dynamics: Exponential convergence and annealing. arXiv preprint arXiv:2202.01009, 2022. Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for overparameterized models using optimal transport. arXiv preprint arXiv:1805.09545, 2018. Yuhao Ding, Junzi Zhang, and Javad Lavaei. On the global convergence of momentumbased policy gradient. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, pages 1910–1934. PMLR, 2022. Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized markov decision processes. In Proceedings of the 36th International Conference on Machine Learning (ICML), 2019. 18

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning (ICML), 2018. Richard Holley and Daniel Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of Statistical Physics, 46(5–6):1159–1194, 1987. Sham M. Kakade. A natural policy gradient. In Advances in Neural Information Processing Systems 14 (NeurIPS 2001), 2001. Sajad Khodadadian, Prakirt Raj Jhunjhunwala, Sushil Mahavir Varma, and Siva Theja Maguluri. On the linear convergence of natural policy gradient algorithm. In 2021 60th IEEE Conference on Decision and Control (CDC), pages 3794–3799, 2021. doi: 10.1109/CDC45484.2021.9682908. Sajad Khodadadian, Prakirt Raj Jhunjhunwala, Sushil Mahavir Varma, and Siva Theja Maguluri. On linear and super-linear convergence of natural policy gradient algorithm. Systems & Control Letters, 164:105214, 2022. doi: 10.1016/j.sysconle.2022.105214. Guanghui Lan. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes. Mathematical Programming, 198 (1):1059–1106, 2023. Guanghui Lan, Yan Li, and Tuo Zhao. Block policy mirror descent. SIAM Journal on Optimization, 33(3):2341–2378, 2023. doi: 10.1137/22M1480409. 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. arXiv preprint arXiv:1509.02971, 2015. Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665–E7671, 2018. Ted Moskovitz et al. Efficient wasserstein natural gradient for policy optimization. arXiv preprint arXiv:2010.05380, 2020. Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In International Conference on Artificial Intelligence and Statistics, pages 9741–9757. PMLR, 2022. David Pfau, Ian Davies, Diana L. Borsa, João Guilherme Madeira Araújo, Brendan Daniel Tracey, and Hado van Hasselt. Wasserstein policy optimization. arXiv preprint arXiv:2505.00663, 2025. John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization. In Proceedings of the 32nd International Conference on Machine Learning (ICML), 2015. 19

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. David Šiška and Yufei Zhang. A note on convergence of wasserstein policy optimization. arXiv preprint arXiv:2605.22622, 2026. Jun Song, Niao He, Lijun Ding, and Chaoyue Zhao. Provably convergent policy optimization via metric-aware trust region methods. Transactions on Machine Learning Research, 2024. Accepted; initial version 2023. Richard S. Sutton, David A. McAllester, Satinder P. Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems 12 (NIPS 1999), pages 1057–1063, 2000. Taiji Suzuki, Denny Wu, and Atsushi Nitanda. Convergence of mean-field langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction. In Advances in Neural Information Processing Systems, 2023. Antonio Terpin, Nicolas Lanzetti, Batuhan Yardim, Florian Dörfler, and Giorgia Ramponi. Trust region policy optimization with optimal transport discrepancies: Duality and algorithm for continuous actions. In Advances in Neural Information Processing Systems (NeurIPS), 2022. Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm control: Software and tasks for continuous control. Software Impacts, 6:100022, 2020. Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Advances in Neural Information Processing Systems, volume 32, 2019. Cédric Villani. Optimal Transport: Old and New, volume 338 of Grundlehren der mathematischen Wissenschaften. Springer, 2009. Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8(3–4):229–256, 1992. Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 23:1–36, 2022. Kakei Yamamoto, Kazusato Oko, Zhuoran Yang, and Taiji Suzuki. Mean field langevin actor-critic: Faster convergence and global optimality beyond lazy learning. In Fortyfirst International Conference on Machine Learning, 2024. Rui Yuan, Simon S. Du, Robert M. Gower, Alessandro Lazaric, and Lin Xiao. Linear convergence of natural policy gradient methods with log-linear policies. In International Conference on Learning Representations (ICLR), 2023. 20

Wenhao Zhan, Shicong Cen, Baihe Huang, Yuxin Chen, Jason D. Lee, and Yuejie Chi. Policy mirror descent for regularized reinforcement learning: A generalized framework with linear convergence. SIAM Journal on Optimization, 33(2):1061–1091, 2023. doi: 10.1137/21M1456789. Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities. In Advances in Neural Information Processing Systems, volume 33, pages 4572–4583, 2020. Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization as wasserstein gradient flows. In International Conference on machine learning, pages 5737–5746. PMLR, 2018. Zhaoyu Zhu, Shuhan Zhang, Rui Gao, and Shuang Li. Wasserstein proximal policy gradient. arXiv preprint arXiv:2603.02576, 2026. Brian D. Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy. PhD thesis, Carnegie Mellon University, 2010.

Appendix A. Auxiliary Results A.1 Entropy-Regularized Performance Difference Lemma Lemma A.1 (Entropy-Regularized Performance Difference Lemma) R (Lan, 2023, Lemma2) Let ρ0 ∈ P(S) be the initial state distribution, and write J(π) := Jρ0 (π) = S V π (s)ρ0 (ds). For any two feasible policies π and π ′ , we have Z Z 1 ′ ′ J(π ′ ) − J(π) = Qπ (s, a) π ′ (a | s) − π(a | s) da − τ Hπ (s) + τ Hπ (s) dπρ0 (ds), 1 − γ S Rd where

π(a | s) log π(a | s) da

H (s) := Rd

is the statewise negative entropy, and ′

dπρ0 (B) := (1 − γ)

∞ X

γ t Pρ,π′ (st ∈ B),

B ∈ B(S),

t=0

is the normalized discounted state occupancy measure of π ′ initialized from ρ. For completeness, we provide the proof here and adapt it to our notation. ′

Proof Let Γπ (ρ0 ) denote the trajectory (st , at , st+1 )t≥0 generated by s0 ∼ ρ0 , at ∼ π ′ (· | st ), and st+1 ∼ P (· | st , at ). By the definition of J, we have "∞ # X J(π ′ ) − J(π) = EΓπ′ (ρ0 ) γ t r̃(st , at ) − τ log π ′ (at | st ) − Es0 ∼ρ0 [V π (s0 )]. t=0

Adding and subtracting V π (st ) inside the summation and rearranging the telescoping value terms gives "∞ # X J(π ′ ) − J(π) = EΓπ′ (ρ0 ) γ t r̃(st , at ) − τ log π ′ (at | st ) + γV π (st+1 ) − V π (st ) t=0 + EΓπ′ (ρ0 ) [V π (s0 )] − Es0 ∼ρ0 [V π (s0 )].

The last two terms cancel. Taking conditional expectation over st+1 given (st , at ) and using Qπ (st , at ) = r̃(st , at ) + γE[V π (st+1 ) | st , at ], we obtain J(π ′ ) − J(π) = EΓπ′ (ρ0 )

"∞ X

# γ t Qπ (st , at ) − V π (st ) − τ log π ′ (at | st )

t=0

For each fixed st , since at ∼ π ′ (· | st ) and Z V π (st ) = Qπ (st , a) − τ log π(a | st ) π(a | st ) da, Rd

we have

E Qπ (st , at ) − V π (st ) − τ log π ′ (at | st ) st Z ′ = Qπ (st , a) π ′ (a | st ) − π(a | st ) da − τ Hπ (st ) + τ Hπ (st ). Rd

Substituting this identity into the trajectory expectation yields "∞ # X Z ′ J(π ′ ) − J(π) = EΓπ′ (ρ0 ) γt Qπ (st , a) π ′ (a | st ) − π(a | st ) da − τ Hπ (st ) + τ Hπ (st ) . Rd

t=0 ′

Finally, by the definition of dπρ0 , the last display is equivalent to Z Z 1 ′ ′ π ′ π′ π J(π ) − J(π) = Q (s, a) π (a | s) − π(a | s) da − τ H (s) + τ H (s) dπρ0 (ds). 1 − γ S Rd This proves the lemma.

A.2 Derivation of the first variation and Wasserstein policy-gradient direction We justify the statewise direction used in Section 2.2. Throughout this subsection, we write J(π) = Jρ0 (π) for the entropy-regularized objective with initial distribution ρ0 . Fix a policy π, and let dπ = dπρ0 be its normalized discounted state occupancy. For each state s, consider a smooth transport perturbation t 7→ πt (· | s) with π0 = π and velocity field u(s, ·), so that ∂t πt (a | s) t=0 = −∇a · π(a | s)u(s, a) . (26) The following proposition identifies the first variation of J along such statewise perturbations. 22

Proposition A.1 (First variation under statewise transport perturbations) Let t 7→ πt be a smooth curve of feasible policies with π0 = π, and write π̇(a | s) := ∂t πt (a | s) t=0 . Then

Z Z

d 1 J(πt ) = dt 1 − γ t=0

Qπ (s, a) − τ log π(a | s) π̇(a | s) da dπ (ds).

(27)

If, in addition, π̇ is generated by the transport perturbation (26), then Z Z d 1 J(πt ) ∇a Qπ (s, a) − τ log π(a | s) , u(s, a) π(a | s) da dπ (ds). (28) = dt 1 − γ S Rd t=0 Proof Apply the entropy-regularized performance difference lemma with π ′ = πt . This gives Z Z 1 π π πt J(πt ) − J(π) = Q (s, a) πt (a | s) − π(a | s) da − τ H (s) + τ H (s) dπt (ds). 1 − γ S Rd Define

Z Ft (s) :=

Qπ (s, a) πt (a | s) − π(a | s) da − τ Hπt (s) + τ Hπ (s).

R Then F0 (s) = 0 for every s. Hence, when differentiating S Ft (s)dπt (ds) at t = 0, the derivative of the occupancy measure does not contribute, and we obtain Z d 1 J(πt ) = Ḟ0 (s) dπ (ds). dt 1−γ S t=0 It remains to compute Ḟ0 (s). Since Z

πt

πt (a | s) log πt (a | s) da,

H (s) = Rd

we have

Because

d πt H (s) = dt t=0 R

1 + log π(a | s) π̇(a | s) da.

Rd π̇(a | s) da = 0, the constant term drops out, and therefore

Z Ḟ0 (s) =

Qπ (s, a) − τ log π(a | s) π̇(a | s) da.

Substituting this expression into the previous display proves (27). Now suppose that π̇ is given by (26). Substituting π̇ = −∇a · (πu) into (27) and using integration by parts in divergence form with respect to the action variable yields Z Z π Q (s, a)−τ log π(a | s) π̇(a | s) da = ∇a Qπ (s, a) − τ log π(a | s) , u(s, a) π(a | s) da. Rd

This proves (28). 23

Thus, the action-space gradient direction associated with the first variation is ∇a Qπ (s, a) − τ log π(a | s) .

(29)

Equivalently, one may view (28) as taking the Wasserstein gradient with respect to the statewise 2-Wasserstein geometry on the conditional action laws, averaged over states using the policy-gradient weight dπ . This is the standard Otto-calculus interpretation of gradients on probability measures (Ambrosio et al., 2008; Villani, 2009); the appearance of dπ is the standard discounted state weighting from the policy-gradient theorem (Sutton et al., 2000), and analogous visitation-weighted averages appear in natural-gradient and trustregion policy optimization (Kakade, 2001; Schulman et al., 2015). Under the full-support convention on ρ, the induced occupancy dπ has full support, so this statewise direction is determined on all states relevant to the sup-norm analysis. Substituting the direction (29) into the continuity equation gives the Wasserstein policygradient flow: ∂t πt (a | s) = −∇a · πt (a | s)∇a Qπt (s, a) − τ log πt (a | s) , which is (WPGF).

Appendix B. Proofs for Section 3 Proof [Proof of Lemma 2.1] Fix t ≥ 0 and s ∈ S. Let Z πt Z := exp(f (a)/τ ) da, f (a) := Q (s, a), A

and define the Gibbs density p(a) :=

exp(f (a)/τ ) , Z

so that p(·) = pπs t (·). For any density µ on A, Z Z Z µ(a) KL(µ∥p) = µ(a) log da = µ(a) log µ(a) da − µ(a) log p(a) da. p(a) A A A Since log p(a) = f (a)/τ − log Z, we have Z Z 1 µ(a) log p(a) da = f (a)µ(a) da − log Z. τ A A Substituting gives Z µ(a) log µ(a) da −

KL(µ∥p) = A

or equivalently, Z

Z f (a)µ(a) da + log Z, A

Z f (a)µ(a) da − τ

1 τ

µ(a) log µ(a) da = τ log Z − τ KL(µ∥p). A

(30)

By definition of T ⋆ (as the supremum over policies applied to the evaluation functional at V πt ), Z Z ⋆ πt πt (T V )(s) = sup Q (s, a)µ(a | s) da − τ µ(a | s) log µ(a | s) da . A

µ(·|s)

Applying (30) with µ(· | s) shows the supremum equals τ log Z and is attained at µ = p. Hence, (T ⋆ V πt )(s) = τ log Z. On the other hand, V πt is the fixed point of T πt , so Z Z πt (a | s) log πt (a | s) da. Qπt (s, a)πt (a | s) da − τ V πt (s) = (T πt V πt )(s) = A

Applying (30) with µ(· | s) = πt (· | s) yields V πt (s) = τ log Z − τ KL(πt (· | s)∥p). Combining the last two displays gives (T ⋆ V πt )(s) − V πt (s) = τ KL(πt (· | s)∥pπs t ), which is nonnegative since KL divergence is nonnegative. Proof [Proof of Lemma 3.1] Fix t ≥ 0 and s ∈ S. The soft policy-evaluation identity gives Z Z πt πt V (s) = Q (s, a)πt (a | s) da − τ πt (a | s) log πt (a | s) da. (31) Rd

Differentiating (31) in t gives Z Z πt Qπt (s, a)∂t πt (a | s) da V̇t (s) = (∂t Q )(s, a)πt (a | s) da + Rd Z Rd −τ ∂t πt (a | s) 1 + log πt (a | s) da. Rd

We first evaluate the term containing ∂t Qπt . Since Z πt Q (s, a) = r̃(s, a) + γ V πt (s′ )P (ds′ | s, a), S

and r̃ and P are time-independent, (∂t Qπt )(s, a) = γ

V̇t (s′ )P (ds′ | s, a).

Therefore, Z

(∂t Qπt )(s, a)πt (a | s) da = γ(P πt V̇t )(s).

(32)

It remains to evaluate the two terms in (32) involving ∂t πt (· | s). Since πt (· | s) is a probability density for every t, Z ∂t πt (a | s) da = 0. Rd

Hence these two terms can be written as Z Qπt (s, a) − τ log πt (a | s) ∂t πt (a | s) da. Rd

By the Gibbs definition of pπs t , pπs t (a) = R

exp(Qπt (s, a)/τ ) , πt Rd exp(Q (s, ā)/τ ) dā

and therefore Qπt (s, a) = τ log pπs t (a) + τ log

exp(Qπt (s, ā)/τ ) dā.

Substituting the two evaluated contributions back into (32), we obtain V̇t (s) = γ(P πt V̇t )(s) + gt (s). Equivalently, (I − γP πt )V̇t = gt , and hence V̇t = (I − γP πt )−1 gt . Since P πt is positivity preserving and γ ∈ (0, 1), (I − γP πt )−1 =

∞ X

γ n (P πt )n

n=0

is also positivity preserving. Because gt (s) ≥ 0 for every s, we get V̇t = gt +

∞ X

γ n (P πt )n gt ≥ gt ≥ 0

n=1

pointwise on S. This proves both the resolvent identity and the pointwise lower bound.

Appendix C. Uniform regularity bounds This appendix verifies the quantitative bounds needed to use the core Bellman proof uniformly along WPGD. Pure measurability of kernels and conditional versions remains covered by the measurability convention in Section 2; Normalizability, Bellman differentiability, value bounds, drift bounds, moments, KL bounds, and LSI are proved below. C.1 Explicit constants and feasible step sizes For Vmax ≥ 0, define Lb (Vmax ) := β + Lr + γLp Vmax , Let U :=

Rmax + τ log Zβ , 1−γ

L⋆ :=

G(Vmax ) := Gr + γGp Vmax .

−Rmax + τ log Zβ , 1−γ

Ē0 :=

2Rmax + τ K0 , 1−γ

and set V̄ := max{1, U, Ē0 − L⋆ + 1}.

(33)

The induced constants are L̄b := Lb (V̄ ),

Ḡ := G(V̄ ),

β 2(Rmax + γ V̄ ) ᾱ := exp − , τ τ

With M0 = sups E∥As0 ∥2 , define 2 Ḡ2 4Ḡ2 M̄∞ (η) := + 2τ d + η, β β β

cη := 1 − e−ᾱτ η .

B 2 := 2β 2 max{M0 M̄∞ (η)} + 2Ḡ2 27

The statewise one-step KL discretization error is δ̄η :=

2 L̄2b d 2 L̄2b B∞ η + η3. 2 6τ

(34)

The auxiliary KL bounds are d βM∞ + log Zβ − log(4πeτ η) , K̄η := max K0 , 0, 2τ 2

2(Rmax + γ V̄ ) . τ The exact stability condition used in Proposition 4.1 is H̄η := K̄η +

0<η≤

1 , 4β

ᾱτ η ≤ 1,

For the simplified theorem statement, let 3Ḡ2 4τ d M̄ := max M0 , 2 + , β β Cδ := and

δ̄η τ ≤ 1. 2 (1 − γ) cη

B̄ 2 := 2β 2 M̄ + 2Ḡ2 ,

L̄2b d L̄2b B̄ 2 + , 2 6τ

1 1 ᾱ(1 − γ)2 η0 := min 1, , , 4β ᾱτ 2Cδ

(35)

(36) .

(37)

For every 0 < η ≤ η0 , Lemma C.6 implies (35), and (34) satisfies δ̄η ≤ Cδ η 2 . C.2 Bellman well-posedness and optimal Gibbs policies Lemma C.1 (Soft optimality contraction and fixed point) Under Assumptions 4.1 and 4.2, with the measurability convention of Section 2, the soft optimality operator T ⋆ maps bounded measurable functions to bounded measurable functions, is monotone, and is a γ-contraction in ∥ · ∥∞ : ∥T ⋆ V − T ⋆ W ∥∞ ≤ γ∥V − W ∥∞ . Consequently, T ⋆ has a unique bounded measurable fixed point V ⋆ . Proof Let Bb (S) denote the bounded measurable functions on S with the sup norm. The measurability part of the mapping property is covered by the convention in Section 2. The quantitative boundedness follows from the bounded-Gaussian representation used below: for bounded V , R r(s, a) + γ V (s′ )P (ds′ | s, a) exp(QV (s, a)/τ ) = Zβ exp ρβ (a), τ and the exponential factor multiplying ρβ is uniformly bounded above and below by positive constants depending only on ∥V ∥∞ . 28

Monotonicity is immediate from V ≤ W ⇒ QV ≤ QW . Since QV (s, a) ≤ QW (s, a) + γ∥V − W ∥∞ for all (s, a), (T ⋆ V )(s) ≤ (T ⋆ W )(s) + γ∥V − W ∥∞ . Interchanging V and W gives ∥T ⋆ V − T ⋆ W ∥∞ ≤ γ∥V − W ∥∞ . The Banach fixed-point theorem on Bb (S) gives a unique bounded measurable fixed point V ⋆ .

Lemma C.2 (Bellman well-posedness and WPG update) Assume the bounded-reward condition in Assumption 4.1(i), the measurability convention in Section 2, and the actionregularity bounds in Assumption 4.2. Let π ∈ Π. Then the following statements hold. (i) The policy-evaluation Bellman operator T π is well-defined on bounded measurable functions, maps them to bounded functions, is a γ-contraction in ∥ · ∥∞ , and has a unique bounded fixed point V π . (ii) For every state s, the action-value function Z β π 2 Q (s, a) = r(s, a) − ∥a∥ + γ V π (s′ )P (ds′ | s, a) 2 S is continuously differentiable in a, and Z

∇a Q (s, a) = ∇a r(s, a) − βa + γ

V π (s′ )∇a p(s′ | s, a)λ(ds′ ).

Moreover, if ∥V π ∥∞ ≤ Vmax , then ∥∇a Qπ (s, a) − ∇a Qπ (s, ā)∥ ≤ (β + Lr + γLp Vmax )∥a − ā∥, and ∇a Qπ (s, a) = −βa + hπ (s, a),

sup ∥hπ (s, a)∥ ≤ Gr + γGp Vmax . s,a

(iii) For every bounded measurable V : S → R, the Gibbs kernel exp(QV (s, a)/τ ) Rd exp(QV (s, ã)/τ )dã

G[V ](a | s) = R

is well-defined, belongs to Π, and attains the variational supremum defining T ⋆ V . (iv) In particular, with V ⋆ fromC.1, the Gibbs kernel π ⋆ := G[V ⋆ ] ⋆

belongs to Π, satisfies V π = V ⋆ , and is an optimal admissible policy. (v) If πk ∈ Π and bk (s, a) := ∇a Qπk (s, a), then the WPG update has the statewise density Z πk+1 (y | s) =

ϕ2τ η (y − a − ηbk (s, a))πk (a | s) da, Rd

where −d/2

ϕ2τ η (z) = (4πτ η)

∥z∥2 exp − . 4τ η

Equivalently, Ask+1 = Ask + ηbk (s, Ask ) +

2τ η ξk+1 ,

Ask ∼ πk (· | s),

ξk+1 ∼ N (0, Id ).

For t > 0, the corresponding one step drift interpolation has density Z s qt (y) = ϕ2τ t (y − a − tbk (s, a))πk (a | s) da, Rd

and its conditional drift is R ust (y) =

Rd bk (s, a)ϕ2τ t (y − a − tbk (s, a))πk (a | s) da qts (y)

for qts (y)dy-a.e. y. Finally, along the WPG sequence, admissibility of the updated densities is obtained from the moment and Gaussian-relative-entropy estimates in Proposition 4.1: under its step-size condition, sup sup KL(πk (· | s)∥ρβ ) < ∞, k≥0 s∈S

and hence all iterates remain in Π. Proof [Proof of Lemma C.2] We prove the quantitative parts of the statement and leave only kernel measurability to the convention in Section 2. First, for each state s, finite Gaussian-relative entropy gives π(· | s) ≪ ρβ ≪ da and a finite second moment by Lemma E.1. Using log ρβ (a) = − log Zβ −

β ∥a∥2 , 2τ

the one-step regularized reward can be written as the finite quantity Z β π 2 r̄ (s) := r(s, a) − ∥a∥ − τ log π(a | s) π(a | s) da 2 Rd Z r(s, a)π(a | s) da + τ log Zβ − τ KL(π(· | s)∥ρβ ). = Rd

Consequently, with Cπ := sups KL(π(· | s)∥ρβ ), −Rmax + τ log Zβ − τ Cπ ≤ r̄π (s) ≤ Rmax + τ log Zβ .

(38)

For bounded V , the policy-evaluation operator is Z Z π π (T V )(s) = r̄ (s) + γ V (s′ )p(s′ | s, a)λ(ds′ )π(a | s) da. Rd

The second term is bounded in absolute value by γ∥V ∥∞ , and (38) gives boundedness of T π V . Moreover, ∥T π V − T π W ∥∞ ≤ γ∥V − W ∥∞ . 30

Thus T π is a γ-contraction on the bounded measurable functions and has a unique bounded fixed point V π . Next, fix a bounded function V and define Z V (s′ )p(s′ | s, a)λ(ds′ ). PV (s, a) := S

We verify the derivative-under-the-integral step. Fix s, a, and a direction h ∈ Rd . For ε ̸= 0, the fundamental theorem of calculus gives, for λ-a.e. s′ , p(s′ | s, a + εh) − p(s′ | s, a) = ε

Z 1

∇a p(s′ | s, a + θεh), h dθ.

Subtracting the candidate directional derivative and integrating, we obtain Z PV (s, a + εh) − PV (s, a) − V (s′ ) ∇a p(s′ | s, a), h λ(ds′ ) ε S Z 1Z ≤ ∥V ∥∞ ∇a p(s′ | s, a + θεh) − ∇a p(s′ | s, a) λ(ds′ ) ∥h∥ dθ 0

1 ≤ ∥V ∥∞ Lp |ε|∥h∥2 . 2 Hence a 7→ PV (s, a) is differentiable and Z ∇a PV (s, a) = V (s′ )∇a p(s′ | s, a)λ(ds′ ). S

The same assumption gives the Lipschitz estimate ∥∇a PV (s, a) − ∇a PV (s, ā)∥ ≤ ∥V ∥∞ Lp ∥a − ā∥, and therefore continuity of the derivative in a. Applying this with V = V π proves the displayed formula for ∇a Qπ . If ∥V π ∥∞ ≤ Vmax , then ∥bπ (s, a) − bπ (s, ā)∥ ≤ β∥a − ā∥ + Lr ∥a − ā∥ Z + γ∥V π ∥∞ ∥∇a p(s′ | s, a) − ∇a p(s′ | s, ā)∥λ(ds′ ) S ≤ β + Lr + γLp Vmax ∥a − ā∥. Also bπ (s, a) = −βa + hπ (s, a),

hπ (s, a) := ∇a r(s, a) + γ

Z S

and the Gr , Gp bounds imply ∥hπ (s, a)∥ ≤ Gr + γGp Vmax . 31

V π (s′ )∇a p(s′ | s, a)λ(ds′ ),

We now turn to the soft optimality operator. For bounded V , set R r(s, a) + γ S V (s′ )P (ds′ | s, a) Rmax + γ∥V ∥∞ ψV (s, a) := , CV := . τ τ Then |ψV (s, a)| ≤ CV , and exp

! R r(s, a) − β2 ∥a∥2 + γ V (s′ )P (ds′ | s, a) = Zβ eψV (s,a) ρβ (a). τ

Thus the normalizer lies in Zβ e−CV , Zβ eCV , so T ⋆ V is finite and bounded. The Gibbs density can be written as G[V ](a | s) = R

eψV (s,a) ρβ (a). eψV (s,ã) ρβ (ã)dã

Consequently, Z KL(G[V ](· | s)∥ρβ ) =

Z ψV (s, a)G[V ](a | s) da − log

eψV (s,a) ρβ (a) da

≤ 2CV . Taking the supremum over s gives G[V ] ∈ Π. Finally, the Gibbs variational identity gives, for every admissible density µ, Z Z Z QV (s, a)µ(a) da − τ µ(a) log µ(a) da = τ log eQV (s,a)/τ da − τ KL(µ∥G[V ](· | s)), so G[V ] attains the one-step supremum. For the WPG update, condition on Ask = a. The conditional distribution of Ask + √ ηbk (s, Ask ) + 2τ η ξk+1 is the Gaussian density φ2τ η (y − a − ηbk (s, a)). Integrating this conditional density against πk (a | s)da gives the displayed formula for πk+1 (y | s), and Fubini’s theorem gives Z Z Z πk+1 (y | s) dy = φ2τ η (y − a − ηbk (s, a)) dy πk (a | s) da = 1. Rd

The fixed-drift interpolation density is obtained by the same conditioning argument with η replaced by t. Since the Gaussian kernel is strictly positive, qts (y) > 0 for every t > 0 and y ∈ Rd . The numerator in the displayed formula for ust is finite because bk (s, ·) has at most linear growth by the Lipschitz estimate above and πk (· | s) has a finite second moment. Bayes’ rule then gives ust (y) = E[bk (s, As0 ) | Ast = y] for qts (y)dy-a.e. y. This is the conditional-drift version used in the fixed-drift Fokker–Planck and entropy-differentiation arguments.

C.2.1 Optimal Gibbs policy Lemma C.3 (Optimal Gibbs policy) Let V ⋆ be the unique bounded fixed point of T ⋆ , and define Z V ⋆ (s′ )P (ds′ | s, a),

Q⋆ (s, a) := r̃(s, a) + γ

and π ⋆ (a | s) := R Then π ⋆ ∈ Π,

exp(Q⋆ (s, a)/τ ) . ⋆ Rd exp(Q (s, ã)/τ ) dã

(39)

⋆

T π V ⋆ = T ⋆V ⋆ = V ⋆, and hence

⋆

V π = V ⋆. Moreover, for every admissible policy π, Vπ ≤V⋆

pointwise on S.

Consequently, π ⋆ is an optimal soft policy within the admissible class, and Z ⋆ J(π ) = V ⋆ (s)ρ0 (ds). S

Proof [Proof of Lemma C.3] Since V ⋆ is bounded and |r| ≤ Rmax , Z r(s, a) + γ V ⋆ (s′ )P (ds′ | s, a) ≤ Rmax + γ∥V ⋆ ∥∞ . S

Write

r(s, a) + γ S V ψs⋆ (a) := Then

and therefore

⋆ (s′ )P (ds′ | s, a)

β Q⋆ (s, a) = − ∥a∥2 + τ ψs⋆ (a), 2 β ⋆ ⋆ 2 exp(Q (s, a)/τ ) = exp(ψs (a)) exp − ∥a∥ . 2τ

Let

Rmax + γ∥V ⋆ ∥∞ . τ Then ∥ψs⋆ ∥∞ ≤ C⋆ for every s. Hence the normalizer in (39) satisfies Z e−C⋆ Zβ ≤ exp(Q⋆ (s, a)/τ ) da ≤ eC⋆ Zβ , C⋆ :=

so it is finite and strictly positive. Moreover, exp(ψs⋆ (a)) ρβ (a). ⋆ Rd exp(ψs (ã))ρβ (ã) dã

π ⋆ (a | s) = R

Thus π ⋆ (· | s) is a bounded perturbation of ρβ . In particular, Z Z ⋆ ⋆ ⋆ exp(ψs⋆ (a))ρβ (a) da ≤ 2C⋆ . ψs (a)π (a | s) da − log KL π (· | s) ∥ ρβ = Rd

Taking the supremum over s gives sup KL π ⋆ (· | s) ∥ ρβ ≤ 2C⋆ < ∞. s∈S

Therefore π ⋆ ∈ Π. The required kernel measurability is part of the convention in Section 2. We now show that π ⋆ attains the soft optimality operator at V ⋆ . For each fixed state s, let Z exp(Q⋆ (s, a)/τ ) da.

Zs⋆ :=

By definition, log π ⋆ (a | s) =

Q⋆ (s, a) − log Zs⋆ . τ

Therefore Q⋆ (s, a) − τ log π ⋆ (a | s) = τ log Zs⋆ for every a. Hence (T

π⋆

⋆

(Q⋆ (s, a) − τ log π ⋆ (a | s)) π ⋆ (a | s) da

V )(s) = Rd

= τ log Zs⋆ . On the other hand, the Gibbs variational formula gives (T ⋆ V ⋆ )(s) = τ log Zs⋆ . Since V ⋆ is the fixed point of T ⋆ , ⋆

T π V ⋆ = T ⋆V ⋆ = V ⋆. ⋆

Because π ⋆ ∈ Π, the policy-evaluation operator T π is a γ-contraction on bounded mea⋆ surable functions and has a unique bounded fixed point, namely V π . Since V ⋆ is also a ⋆ bounded fixed point of T π , uniqueness implies ⋆

V π = V ⋆. Finally, for any π ∈ Π and any bounded measurable V , T π V ≤ T ⋆V pointwise, because T ⋆ is the supremum over admissible action densities in the one-step soft Bellman expression. Applying this with V = V ⋆ gives T π V ⋆ ≤ T ⋆V ⋆ = V ⋆. 34

By monotonicity of T π , (T π )n V ⋆ ≤ V ⋆

∀n ≥ 1.

Since T π is a γ-contraction, (T π )n V ⋆ converges uniformly to its unique bounded fixed point V π . Passing to the limit yields V π ≤ V ⋆. ⋆

Thus π ⋆ is optimal among admissible policies. Since V π = V ⋆ , Z Z π⋆ ⋆ V ⋆ (s)ρ0 (ds). V (s)ρ0 (ds) = J(π ) = S

This completes the proof.

C.3 Value bounds and Bellman drift estimates For every π ∈ Π, the one-step regularized reward obeys Z π r̄ (s) := r̃(s, a) − τ log π(a | s) π(a | s) da d ZR = r(s, a)π(da | s) + τ log Zβ − τ KL(π(· | s)∥ρβ ).

(40)

Consequently, under Assumption 4.1, −Rmax + τ log Zβ − τ Cπ ≤ r̄π (s) ≤ Rmax + τ log Zβ ,

Cπ := sup KL(π(· | s)∥ρβ ). s

This coercive identity is the basic reason admissibility relative to ρβ makes Bellman evaluation finite on the unbounded action space. Lemma C.4 (Basic value bounds) Let U :=

Rmax + τ log Zβ , 1−γ

L⋆ :=

−Rmax + τ log Zβ , 1−γ

Ē0 :=

2Rmax + τ K0 . 1−γ

Under Assumption 4.1, every admissible policy π ∈ Π satisfies V π (s) ≤ U

∀s ∈ S.

The optimal value satisfies L⋆ ≤ V ⋆ (s) ≤ U

∀s ∈ S,

and the initial policy satisfies V π0 (s) ≥

−Rmax + τ log Zβ − τ K0 , 1−γ 35

∥V ⋆ − V π0 ∥∞ ≤ Ē0 .

Proof [Proof of Lemma C.4] For π ∈ Π, write Cπ = sups KL(π(· | s)∥ρβ ). The identity (40) gives, for every state s, −Rmax + τ log Zβ − τ Cπ ≤ r̄π (s) ≤ Rmax + τ log Zβ . Thus T π is an affine discounted Bellman operator with uniformly bounded one-step reward. For the constant function U = (Rmax + τ log Zβ )/(1 − γ), T π U ≤ Rmax + τ log Zβ + γU = U. By monotonicity of T π and contraction of policy evaluation, (T π )n U → V π uniformly and hence V π ≤ U . For the lower bound on V ⋆ , use the stationary Gaussian policy π β (da | s) = ρβ (a)da. For this policy, Z Z r̃(s, a) − τ log ρβ (a) ρβ (a)da = r(s, a)ρβ (a)da + τ log Zβ ≥ −Rmax + τ log Zβ . Rd β

Therefore V π ≥ L⋆ . Since T ⋆ V ≥ T π V for every bounded V , monotone comparison β of the discounted fixed points gives V ⋆ ≥ V π ≥ L⋆ . The upper bound V ⋆ ≤ U follows similarly from the log-partition form of T ⋆ : for the constant function U , Z ⋆ (Rmax +γU )/τ −β∥a∥2 /(2τ ) (T U )(s) ≤ τ log e e da = Rmax + γU + τ log Zβ = U, and iterating T ⋆ from U yields V ⋆ ≤ U . Finally, Assumption 4.1(ii) and (40) imply r̄π0 (s) ≥ −Rmax + τ log Zβ − τ K0 . Applying the same monotone-comparison argument to T π0 gives V π0 (s) ≥

−Rmax + τ log Zβ − τ K0 . 1−γ

Since Lemma C.3 gives V π0 ≤ V ⋆ , the preceding initial lower bound and V ⋆ ≤ U yield ∥V ⋆ − V π0 ∥∞ ≤

−Rmax + τ log Zβ − τ K0 Rmax + τ log Zβ − = Ē0 . 1−γ 1−γ

Proof [Proof of Lemma 4.4] Lemma C.2 proves the Bellman action-gradient formula Z bk (s, a) = ∇a r(s, a) − βa + γ V πk (s′ )∇a p(s′ | s, a)λ(ds′ ). S

Set

V πk (s′ )∇a p(s′ | s, a)λ(ds′ ).

hk (s, a) := ∇a r(s, a) + γ S

Then bk (s, a) = −βa + hk (s, a). If ∥V πk ∥∞ ≤ Vmax , Assumption 4.2 gives Z ∥∇a p(s′ | s, a)∥λ(ds′ ) ≤ Gr + γGp Vmax . ∥hk (s, a)∥ ≤ Gr + γVmax S

This proves the dissipative decomposition and (19). For the Lipschitz bound, subtract the displayed gradient formula at a and ā and use the reward-gradient Lipschitz bound and the L1 (λ) Lipschitz bound on ∇a p: ∥bk (s, a) − bk (s, ā)∥ ≤ β∥a − ā∥ + Lr ∥a − ā∥ Z πk ∥∇a p(s′ | s, a) − ∇a p(s′ | s, ā)∥λ(ds′ ) + γ∥V ∥∞ S

≤ (β + Lr + γLp Vmax )∥a − ā∥.

C.4 Moment bounds, step sizes, and uniform stability Lemma C.5 (Second-moment bound under dissipativity) Suppose that for all 0 ≤ k ≤ K − 1, Lemma 4.4 holds with G(Vmax ) = G. Assume M0 := sups∈S E∥As0 ∥2 < ∞ and 1 0 < η ≤ 4β . Then (41) sup sup E∥Ask ∥2 ≤ max{M0 , M∞ (η)}, 0≤k≤K s∈S

where one may take M∞ (η) :=

4G2 2 G2 + 2τ d + η. β β β

Proof [Proof of Lemma C.5] Fix s ∈ S and write Ak := Ask for brevity. Let mk := E∥Ak ∥2 ,

0 ≤ k ≤ K.

Recall the update Ak+1 = Ak + η bk (s, Ak ) +

p 2τ η ξk+1 ,

ξk+1 ∼ N (0, Id ),

with ξk+1 independent of Ak . Taking squared norms and expectations gives mk+1 = E Ak + ηbk (s, Ak ) + 2

= E Ak + ηbk (s, Ak )

2τ η ξk+1

+ 2τ η E∥ξk+1 ∥2

because the cross term vanishes by independence and E[ξk+1 ] = 0. Since E∥ξk+1 ∥2 = d, we have 2 mk+1 = E Ak + ηbk (s, Ak ) + 2τ ηd. (42) Expanding the first term yields 2

E Ak + ηbk (s, Ak )

= mk + 2ηE⟨Ak , bk (s, Ak )⟩ + η 2 E∥bk (s, Ak )∥2 . 37

(43)

By the assumed dissipativity bound (Lemma 4.4 with G(Vmax ) = G), for each 0 ≤ k ≤ K − 1 we may write bk (s, a) = −βa + hk (s, a),

sup ∥hk (s, a)∥ ≤ G. s,a

Hence, for any a ∈ Rd , β G2 ⟨a, bk (s, a)⟩ = −β∥a∥2 + ⟨a, hk (s, a)⟩ ≤ −β∥a∥2 + G∥a∥ ≤ − ∥a∥2 + , 2 2β 2

where the last step uses Young’s inequality G∥a∥ ≤ β2 ∥a∥2 + G 2β . Thus β G2 E⟨Ak , bk (s, Ak )⟩ ≤ − mk + . 2 2β

(44)

Moreover, using (x + y)2 ≤ 2x2 + 2y 2 and ∥hk ∥ ≤ G, ∥bk (s, a)∥2 = ∥ − βa + hk (s, a)∥2 ≤ 2β 2 ∥a∥2 + 2G2 , so E∥bk (s, Ak )∥2 ≤ 2β 2 mk + 2G2 .

(45)

Plugging (44) and (45) into (42)-(43) gives, for 0 ≤ k ≤ K − 1, β G2 mk+1 ≤ mk + 2η − mk + + η 2 2β 2 mk + 2G2 + 2τ ηd 2 2β G2 = 1 − βη + 2β 2 η 2 mk + η + 2τ d + 2G2 η 2 . β 1 Under the step-size condition η ≤ 4β ,

2β 2 η 2 ≤

βη 2

1 − βη + 2β 2 η 2 ≤ 1 −

=⇒

Define ρ := 1 − βη 2 ∈ (0, 1) and ζη := η Then we have the affine recursion

G2 β + 2τ d

mk+1 ≤ ρ mk + ζη ,

βη . 2

+ 2G2 η 2 .

0 ≤ k ≤ K − 1.

Iterating (46) yields k

mk ≤ ρ m0 + ζη

k−1 X

ρj = ρk m0 + ζη

j=0 G2

2 2

ζη 1 − ρk = ρk m0 + (1 − ρk ) . 1−ρ 1−ρ

4G2 η( β + 2τ d) + 2G η ζη 2 G2 = = + 2τ d + η =: M∞ (η). 1−ρ βη/2 β β β 38

(46)

Therefore mk ≤ ρk m0 + (1 − ρk )M∞ (η) ≤ max{m0 , M∞ (η)}

∀ 0 ≤ k ≤ K.

Finally, taking sup0≤k≤K and then sups∈S , and using m0 = E∥As0 ∥2 ≤ M0 , we obtain sup sup E∥Ask ∥2 ≤ max{M0 , M∞ (η)},

0≤k≤K s∈S

which is (41).

Lemma C.6 (Feasibility of the stability step-size condition) Let the constants in (33)(37) be defined as above. Then every 0 < η ≤ η0 satisfies the exact stability condition (35). Moreover, δ̄η ≤ Cδ η 2 . Proof [Proof of Lemma C.6] All constants in (37) are finite and positive under the standing assumptions, so η0 > 0. Fix 0 < η ≤ η0 . The first two inequalities in (35) follow from η≤

1 , 4β

η≤

1 . ᾱτ

Since η ≤ 1/(4β), 4Ḡ2 3Ḡ2 4τ d 2 Ḡ2 + 2τ d + η≤ 2 + . M̄∞ (η) = β β β β β n o Therefore max M0 , M̄∞ (η) ≤ M̄ and B ≤ B̄. Substituting this into (34) and using η ≤ 1 gives 2 L̄2b d 2 L̄2b B 2 3 L̄b d L̄2b B̄ 2 δ̄η = η + η ≤ + η 2 = Cδ η 2 . 2 6τ 2 6τ Finally, since ᾱτ η ≤ 1 and 1 − e−x ≥ x/2 for x ∈ [0, 1], cη = 1 − e−ᾱτ η ≥ Thus

ᾱτ η . 2

δ̄η τ τ Cδ η 2 2Cδ ≤ = η ≤ 1, 2 2 (1 − γ) cη (1 − γ) ᾱτ η/2 ᾱ(1 − γ)2

where the last inequality uses the definition of η0 . Hence (35) holds. Proof [Proof of Proposition 4.1] By Lemma C.6, every 0 < η ≤ η0 satisfies the exact stability condition (35). The proof below uses this exact condition. For any index k for which πk is admissible, write bk (s, a) := ∇a Qπk (s, a). 39

For such k and each s ∈ S, let pπs k (a) = R

exp(Qπk (s, a)/τ ) . πk Rd exp(Q (s, ã)/τ )dã

Define Rk (s) := (T ⋆ V πk )(s) − V πk (s),

Wk (s) := V ⋆ (s) − V πk (s),

Ek := ∥Wk ∥∞ .

By Lemma 2.1, Rk (s) = τ KL(πk (· | s)∥pπs k ). Moreover, by Lemma C.1, T ⋆ is monotone and a γ-contraction; since V ⋆ ≥ V πk pointwise, 0 ≤ V ⋆ (s) − (T ⋆ V πk )(s) = (T ⋆ V ⋆ )(s) − (T ⋆ V πk )(s) ≤ γEk ,

∀s ∈ S.

(47)

For each m ≥ 0, let Im denote the assertion that, for every 0 ≤ j ≤ m, πj ∈ Π,

∥V πj ∥∞ ≤ V̄ .

We prove Im for all m by induction. The admissibility part ensures that Lemma C.2 applies recursively to the Bellman objects and to the WPG push-forward density; measurability is handled throughout by the convention in Section 2. Base case. The admissibility of π0 follows from Assumption 4.1(ii), while kernel measurability is covered by the convention in Section 2. By Lemma C.4, sups V π0 (s) ≤ U ≤ V̄ . By Lemma C.4, V π0 (s) ≥ V ⋆ (s) − E0 ≥ L⋆ − E0 . By Lemma C.4, E0 ≤ Ē0 . The definition (33) implies V̄ ≥ Ē0 − L⋆ + 1, and hence L⋆ − E0 ≥ L⋆ − Ē0 ≥ −V̄ . Thus ∥V π0 ∥∞ ≤ V̄ , so I0 holds. Induction step. Assume Ik holds. Then Lemma C.2 gives the Bellman objects V πj , Qπj , bj for 0 ≤ j ≤ k and the explicit WPG push-forward density for πk+1 . Lemma 4.4 implies that, for all 0 ≤ j ≤ k, ∥bj (s, a) − bj (s, ā)∥ ≤ L̄b ∥a − ā∥,

bj (s, a) = −βa + hj (s, a),

sup ∥hj (s, a)∥ ≤ Ḡ. s,a

Applying Lemma C.5 on the finite horizon {0, 1, . . . , k + 1} with G = Ḡ yields sup

sup E∥Asj ∥2 ≤ max{M0 , M̄∞ (η)} ≤ M̄ .

(48)

0≤j≤k+1 s∈S

Consequently, for every 0 ≤ j ≤ k, sup E∥bj (s, Asj )∥2 ≤ 2β 2 max{M0 , M̄∞ (η)} + 2Ḡ2 ≤ B̄ 2 . s∈S

At the current step k, Lemma 4.3 gives that pπs k satisfies an LSI with constant at least ᾱ. The Gibbs definition and the Bellman-gradient formula in Lemma C.2 also give pπs k > 0, log pπs k ∈ 40

C 1 (Rd ), and the score identity τ ∇a log pπs k = bk (s, ·). We next verify the finite-secondmoment and finite-KL hypotheses in Lemma 4.5 quantitatively, and the admissibility/finiteKL hypotheses needed for Lemma 4.6. For j = 0, Assumption 4.1(ii) gives sup KL(π0 (· | s)∥ρβ ) ≤ K0 ≤ K̄η . s

For 1 ≤ j ≤ k + 1, the update writes p s Asj = Xj−1 + 2τ ηξj ,

s Xj−1 := Asj−1 + ηbj−1 (s, Asj−1 ),

s . Since (48) gives E∥As ∥2 ≤ M , Lemma E.3 with with ξj ∼ N (0, Id ) independent of Xj−1 ∞ j 2 σ = 2τ η yields, for 1 ≤ j ≤ k + 1,

KL(πj (· | s)∥ρβ ) ≤

βM∞ d + log Zβ − log(4πeτ η) ≤ K̄η . 2τ 2

(49)

Combining the initialization bound for j = 0 with the smoothing bound for j ≥ 1 gives the uniform finite-horizon estimate sup

sup KL(πj (· | s)∥ρβ ) ≤ K̄η .

(50)

0≤j≤k+1 s∈S

Thus the Gaussian-KL bound holds for all policies generated up to time k + 1, and in particular πj ∈ Π for all j ≤ k + 1. The remaining kernel measurability is covered by the convention in Section 2, so the admissibility part of the induction is closed. Under the value bound, pπs k is a bounded perturbation of ρβ : R r(s, a) + γ V πk (s′ )P (ds′ | s, a) eψk,s (a) πk ps (a) = R ψ ρβ (a), ψk,s (a) := . τ e k,s dρβ Since |r| ≤ Rmax and ∥V πk ∥∞ ≤ V̄ , ∥ψk,s ∥∞ ≤ (Rmax + γ V̄ )/τ . The bounded-perturbation part of Lemma E.3, applied with the j = k case of (50), gives KL(πk (· | s)∥pπs k ) ≤ K̄η +

2(Rmax + γ V̄ ) = H̄η < ∞. τ

(51)

The finite-second-moment hypothesis of Lemma 4.5 follows from (48); the drift secondmoment hypothesis follows from sup E∥bk (s, Ask )∥2 ≤ 2β 2 max{M0 , M̄∞ (η)} + 2Ḡ2 = B̄ 2 . s∈S

Applying Lemma 4.5 with α = ᾱ, Lb = L̄b , and B 2 = B̄ 2 gives, for every s, KL(πk+1 (· | s)∥pπs k ) ≤ e−ᾱτ η KL(πk (· | s)∥pπs k ) + δ̄η .

(52)

The right-hand side of (52) is finite by (51), and hence KL(πk+1 (· | s)∥pπs k ) < ∞. Thus the admissibility and finite-KL hypotheses of Lemma 4.2 are verified for (πk , πk+1 ). Applying Lemma 4.6 to (52) with cη := 1 − e−ᾱτ η and δ = δ̄η gives gk (s) ≥ cη Rk (s) − τ δ̄η . 41

In particular, since Rk ≥ 0, gk (s) ≥ −τ δ̄η and ∥gk− ∥∞ ≤ τ δ̄η . The one-step resolvent identity gives X γ n (P πk+1 )n gk . V πk+1 − V πk = (I − γP πk+1 )−1 gk = n≥0

For n ≥ 1, (P πk+1 )n gk (s) ≥ −∥gk− ∥∞ , so X γ n ∥gk− ∥∞ (V πk+1 − V πk )(s) ≥ gk (s) − n≥1

≥ cη Rk (s) − τ δ̄η −

γ τ τ δ̄η = cη Rk (s) − δ̄η . 1−γ 1−γ

Equivalently, Wk+1 (s) ≤ Wk (s) − cη Rk (s) +

τ δ̄η . 1−γ

(53)

Using (47), Rk (s) = Wk (s) − V ⋆ (s) − (T ⋆ V πk )(s) ≥ Wk (s) − γEk . Combining this bound with (53) yields Wk+1 (s) ≤ (1 − cη )Wk (s) + cη γEk +

τ τ δ̄η ≤ 1 − (1 − γ)cη Ek + δ̄η . 1−γ 1−γ

Taking the supremum over s gives Ek+1 ≤ κη Ek +

τ δ̄η , 1−γ

κη := 1 − (1 − γ)cη .

(54)

Iterating (54) up to k + 1 gives k

Ek+1 ≤ κk+1 η E0 +

X δ̄η τ τ δ̄η ≤ Ē0 + 1, κjη ≤ E0 + 1−γ (1 − γ)2 cη j=0

where the last inequality is the stability step-size condition. Finally, Lemma C.4 gives V πk+1 (s) ≤ U ≤ V̄ , while Lemma C.4 gives V πk+1 (s) ≥ V ⋆ (s) − Ek+1 ≥ L⋆ − Ē0 − 1 ≥ −V̄ , again by (33). Therefore ∥V πk+1 ∥∞ ≤ V̄ . Together with the admissibility conclusion above, this proves Ik+1 and closes the induction. It remains to record the uniform consequences. The value bound just proved implies the drift Lipschitz bound and dissipative decomposition for every k via Lemma 4.4. Applying Lemma C.5 on an arbitrary finite horizon and then letting that horizon vary gives sups E∥Ask ∥2 ≤ max{M0 , M̄∞ (η)} for every k. The drift second-moment bound follows from ∥bk (s, a)∥2 ≤ 2β 2 ∥a∥2 + 2Ḡ2 . The quantitative Gaussian-KL bound (24) follows by applying (50) on arbitrary finite horizons; the k = 0 case is supplied by Assumption 4.1(ii), while every k ≥ 1 case is supplied by the Gaussian-smoothing estimate (49). The kernel measurability assertion follows from the convention in Section 2. Once the value bound is 42

closed for all iterations, the bounded-perturbation argument above applies to every pπs k and gives (25). Finally, Lemma 4.3 with Vmax = V̄ gives the uniform LSI constant ᾱ for every pπs k . Proof [Proof of Lemma 4.3] Fix s ∈ S. By the definition of pπs k , write R r(s, a) + γ S V πk (s′ )P (ds′ | s, a) β πk 2 ps (a) ∝ exp exp − ∥a∥ . τ 2τ Thus pπs k is a bounded perturbation of the Gaussian ρβ (a) ∝ exp(−β∥a∥2 /(2τ )). If ∥V πk ∥∞ ≤ Vmax , then the perturbation R r(s, a) + γ S V πk (s′ )P (ds′ | s, a) ψk,s (a) := τ satisfies

Rmax + γVmax 2(Rmax + γVmax ) , osc(ψk,s ) ≤ . τ τ Since ρβ satisfies an LSI with constant β/τ , the Holley-Stroock bounded perturbation lemma implies that pπs k satisfies an LSI with constant at least β 2(Rmax + γVmax ) exp − = α(Vmax ). τ τ |ψk,s (a)| ≤

This proves the claim.

Appendix D. Core Bellman Contraction Proof This appendix proves the Bellman part of the argument. These are the nonstandard identities that replace the flat-convex gap-conversion step in mean-field Langevin analysis. D.1 Bellman residual, residual-to-gap, and resolvent identities Proof [Proof of Lemma 4.1] Recall Wk (s) = V ⋆ (s) − V πk (s),

Ek = ∥Wk ∥∞ ,

Rk (s) = (T ⋆ V πk )(s) − V πk (s).

Decompose Wk by adding and subtracting (T ⋆ V πk )(s): Wk (s) = V ⋆ (s)−(T ⋆ V πk )(s) + (T ⋆ V πk )(s)−V πk (s) = V ⋆ −(T ⋆ V πk ) (s)+Rk (s). (55) Since V ⋆ = T ⋆ V ⋆ and T ⋆ is a γ-contraction in ∥ · ∥∞ by Lemma C.1, ∥V ⋆ − T ⋆ V πk ∥∞ = ∥T ⋆ V ⋆ − T ⋆ V πk ∥∞ ≤ γ∥V ⋆ − V πk ∥∞ = γEk . In particular, for every s ∈ S, V ⋆ (s) − (T ⋆ V πk )(s) ≤ γEk . 43

(56)

Using (55) and (56), Rk (s) = Wk (s) − V ⋆ − (T ⋆ V πk ) (s) ≥ Wk (s) − γEk .

Proof [Proof of Lemma 4.2] Recall the policy evaluation operator T π (2) and the induced state kernel P π (ds′ | s). For a fixed policy π, define the regularized immediate reward Z π r̃(s, a) − τ log π(a | s) π(a | s) da. r̄ (s) := A

Then we can rewrite T π as the affine map (T π V )(s) = r̄π (s) + γ(P π V )(s). Since V πk+1 is the fixed point of T πk+1 , we have V πk+1 = T πk+1 V πk+1 . Subtract V πk from both sides and add and subtract T πk+1 V πk : V πk+1 − V πk = T πk+1 V πk+1 − T πk+1 V πk + T πk+1 V πk − V πk = γP πk+1 (V πk+1 − V πk ) + gk , which rearranges to the claimed resolvent identity (I − γP πk+1 )(V πk+1 − V πk ) = gk . This proves (16). Now compute (T πk+1 V πk )(s) using the definition of T π and the definition of Qπk : Z πk+1 πk (T V )(s) = Qπk (s, a) − τ log πk+1 (a | s) πk+1 (a | s)da A

Similarly, since V πk = T πk V πk , V

πk

(s) =

Qπk (s, a) − τ log πk (a | s) πk (a | s)da.

Using the standard Gibbs identity Z Z f (a)µ(a) da − τ µ(a) log µ(a) da = τ log Z − τ KL(µ∥p), A

we have that (T πk+1 V πk )(s) = τ log Z − τ KL(πk+1 (· | s)∥pπs k ), V πk (s) = τ log Z − τ KL(πk (· | s)∥pπs k ). Subtracting gives gk (s) = (T πk+1 V πk )(s) − V πk (s) = τ KL(πk (· | s)∥pπs k ) − KL(πk+1 (· | s)∥pπs k ) , which is exactly (17).

D.2 One-step KL contraction with Discretization Error Lemma D.1 (Fokker–Planck with conditional drift) Let A0 ∼ πk (· | s) and let (Bt )t≥0 be standard Brownian motion in Rd , independent of A0 . Let b : Rd → Rd be Borel and assume E∥b(A0 )∥ < ∞. For t ∈ [0, η] define √ At := A0 + t b(A0 ) + 2τ Bt , µt := Law(At ). Define the conditional drift ut (a) := E[b(A0 ) | At = a]. Then, for t > 0, µt satisfies the Fokker–Planck equation ∂t µt = −∇ · (ut µt ) + τ ∆µt in the sense of distributions. Proof [Proof of Lemma D.1] Condition on A0 = a0 . Then √ At = a0 + t b(a0 ) + 2τ Bt , so the conditional law of At has density pt (a | a0 ) = φ2τ t a − a0 − tb(a0 ) , where φ2τ t is the Gaussian density with covariance 2τ t Id . Equivalently, pt (· | a0 ) solves the standard Fokker–Planck equation ∂t pt (· | a0 ) = −∇ · b(a0 ) pt (· | a0 ) + τ ∆pt (· | a0 ). Now the marginal density of µt is the mixture Z µt (a) = pt (a | a0 ) πk (da0 | s). Integrating the above Fokker–Planck equation over a0 gives Z ∂t µt (a) = −∇ · b(a0 ) pt (a | a0 ) πk (da0 | s) + τ ∆µt (a). Define the flux

Z jt (a) :=

b(a0 ) pt (a | a0 ) πk (da0 | s).

By Bayes’ rule, for µt -a.e. a, R ut (a) = E[b(A0 ) | At = a] =

b(a0 ) pt (a | a0 ) πk (da0 | s) jt (a) R . = µt (a) pt (a | a0 ) πk (da0 | s)

Hence jt (a) = ut (a) µt (a), and substituting this into the previous display yields ∂t µt = −∇ · (ut µt ) + τ ∆µt , 45

as claimed. Proof [Proof of Lemma 4.5] Fix k ≥ 0 and s ∈ S. For brevity write p(a) := pπs k (a).

b(a) := bk (s, a),

Let A0 ∼ πk (· | s) and let {Bt }t≥0 be standard Brownian motion in Rd , independent of A0 . For t ∈ [0, η] define the one-step interpolation √ At := A0 + t b(A0 ) + 2τ Bt , µt := Law(At ). (57) Then µ0 = πk (· | s) and µη = πk+1 (· | s). We denote the density of µt again by µt (a) and define ut (a) := E[b(A0 ) | At = a]. By Lemma D.1, µt solves ∂t µt = −∇ · (ut µt ) + τ ∆µt . Set

Z K(t) := KL(µt ∥p) =

µt (a) log Rd

µt (a) da. p(a)

Using the Fokker–Planck equation and differentiating under the integral sign, Z µt (a) d ′ µt (a) log da K (t) = dt d p(a) Z R Z µt (a) = log ∂t µt (a) da + ∂t µt (a) da. p(a) Rd Rd R Since Rd ∂t µt (a) da = 0, this becomes Z µt (a) K ′ (t) = log [−∇ · (ut (a)µt (a)) + τ ∆µt (a)] da. p(a) Rd Integrating by parts gives Z Z µt µt ′ ∇ log (a), ∇µt (a) da K (t) = ∇ log (a), ut (a) µt (a) da − τ p p Rd Rd Z Z µt µt = ∇ log (a), ut (a) µt (a) da − τ ∇ log (a), ∇ log µt (a) µt (a) da. p p Rd Rd By the Gibbs-score identity, τ ∇ log p = b. Therefore ∇ log µt = ∇ log

µt 1 µt + ∇ log p = ∇ log + b. p p τ

Substituting this identity into the preceding display yields h i µt K ′ (t) = Eµt ut (A) − b(A), ∇ log (A) − τ I(µt ∥p). p 46

(58)

Young’s inequality gives h i µt τ 1 Eµt ut (A) − b(A), ∇ log (A) ≤ Eµt ∥ut (A) − b(A)∥2 + I(µt ∥p). p 2τ 2 Combining this with (58), τ 1 K ′ (t) ≤ − I(µt ∥p) + Eµt ∥ut (A) − b(A)∥2 . 2 2τ By the LSI for p, I(µt ∥p) ≥ 2αKL(µt ∥p) = 2αK(t). Hence

1 Eµ ∥ut (A) − b(A)∥2 . 2τ t It remains to bound the interpolation error. By conditional Jensen’s inequality, K ′ (t) ≤ −ατ K(t) +

Eµt ∥ut (A) − b(A)∥2 = E ∥E[b(A0 ) | At ] − b(At )∥2 = E ∥E[b(A0 ) − b(At ) | At ]∥2 ≤ E E ∥b(A0 ) − b(At )∥2 | At = E∥b(A0 ) − b(At )∥2 . Using the Lb -Lipschitzness of b, Eµt ∥ut (A) − b(A)∥2 ≤ L2b E∥A0 − At ∥2 . From (57), At − A0 = t b(A0 ) +

√

2τ Bt .

Since Bt is independent of A0 and has mean zero, E∥At − A0 ∥2 = t2 E∥b(A0 )∥2 + 2τ d t ≤ B 2 t2 + 2τ d t. Thus Eµt ∥ut (A) − b(A)∥2 ≤ L2b B 2 t2 + 2τ L2b d t. Substituting this bound into (59), we get K ′ (t) ≤ −ατ K(t) + L2b d t +

L2b B 2 2 t . 2τ

By Gronwall’s inequality, Z η

L2b B 2 2 2 K(η) ≤ e K(0) + e Lb d u + u du 2τ 0 L2 d L2 B 2 ≤ e−ατ η K(0) + b η 2 + b η 3 . 2 6τ −ατ η

−ατ (η−u)

Recalling that K(0) = KL(πk (· | s)∥pπs k ),

K(η) = KL(πk+1 (· | s)∥pπs k ), 47

(59)

we obtain KL(πk+1 (· | s)∥pπs k ) ≤ e−ατ η KL(πk (· | s)∥pπs k ) +

L2b d 2 L2b B 2 3 η + η . 2 6τ

This proves (20).

D.3 From statewise KL contraction to Bellman improvement Proof [Proof of Lemma 4.6] By Lemma 4.2, gk (s) = τ KL(πk (· | s)∥pπs k ) − KL(πk+1 (· | s)∥pπs k ) . Using (21), gk (s) ≥ τ KL(πk (· | s)∥pπs k ) − (1 − c)KL(πk (· | s)∥pπs k ) − δ = c τ KL(πk (· | s)∥pπs k ) − τ δ = cRk (s) − τ δ. This proves (22). The final sentence follows by substituting the KL contraction from Lemma 4.5.

D.4 Proof of the main theorem Proof [Proof of Theorem 4.1] Define Wk (s) := V ⋆ (s) − V πk (s),

Ek := ∥Wk ∥∞ ,

Rk (s) := (T ⋆ V πk )(s) − V πk (s).

By Proposition 4.1, the constants L̄b , M̄ , B̄ 2 , ᾱ, δ̄η , and cη are valid uniformly over all iterates. The Gibbs targets are strictly positive bounded perturbations of ρβ , the score identity (15) holds, and the finite-second-moment, drift-moment, admissibility, and KL(πk (· | s)∥pπs k ) < ∞ hypotheses needed for Lemmas 4.5 and 4.2 hold at every step. The one-step contraction below follows from the same argument used in the proof of Proposition 4.1; for completeness and readability, we spell out the derivation. Lemma 4.5, applied with (α, Lb , B 2 ) = (ᾱ, L̄b , B̄ 2 ), gives KL(πk+1 (· | s)∥pπs k ) ≤ e−ᾱτ η KL(πk (· | s)∥pπs k ) + δ̄η . Applying Lemma 4.6 with cη = 1 − e−ᾱτ η yields gk (s) ≥ cη Rk (s) − τ δ̄η . Since Rk ≥ 0, this also gives gk (s) ≥ −τ δ̄η . The one-step resolvent identity gives V πk+1 − V πk = (I − γP πk+1 )−1 gk =

X n≥0

γ n (P πk+1 )n gk ,

and hence (V πk+1 − V πk )(s) ≥ cη Rk (s) −

τ δ̄η . 1−γ

Wk+1 (s) ≤ Wk (s) − cη Rk (s) +

τ δ̄η . 1−γ

Therefore

(60)

By Lemma 4.1, Rk (s) ≥ Wk (s) − γEk . Substituting this into (60) and taking the supremum over s gives τ δ̄η , κη := 1 − (1 − γ)cη . Ek+1 ≤ κη Ek + 1−γ Iterating, Ek ≤ κkη E0 +

δ̄η τ . 2 (1 − γ) cη

Since 0 < η ≤ η0 , we have ᾱτ η ≤ 1, so cη = 1−e−ᾱτ η ≥ ᾱτ η/2, and κkη ≤ exp(−(1−γ)cη k) ≤ exp(− 21 ᾱτ (1 − γ)ηk). Also, by (36) and (37), δ̄η ≤ Cδ η 2 . Thus ∥V ⋆ − V πk ∥∞ = Ek ≤ exp − 12 ᾱτ (1 − γ)ηk ∥V ⋆ − V π0 ∥∞ +

2Cδ η. ᾱ(1 − γ)2

Finally, Lemma C.3 gives V πk ≤ V ⋆ , and hence Z J(π ⋆ ) − J(πk ) = (V ⋆ (s) − V πk (s))ρ0 (ds) ≤ ∥V ⋆ − V πk ∥∞ . S

Appendix E. Standard analytic tools This appendix collects standard analytic inputs used in the fixed-drift Langevin estimate and the uniform bounds. These tools are stated separately so that the paper’s main proof remains focused on the Bellman residual/resolvent mechanism. E.1 Gaussian entropy, smoothing, and bounded perturbations Lemma E.1 (Finite Gaussian-relative entropy gives moment) Let µ be R a second d 2 a probability measure on R . If KL(µ∥ρβ ) < ∞, then ∥a∥ µ(da) < ∞. More quantitatively, for every c ∈ (0, β/(2τ )), Z Z 2 2 c ∥a∥ µ(da) ≤ KL(µ∥ρβ ) + log ec∥a∥ ρβ (a) da < ∞. Rd

Proof We use the entropy variational inequality: for any probability measures µ, ν and any measurable function f such that the right-hand side is well-defined, Z Z f dµ ≤ KL(µ∥ν) + log ef dν. 49

Applying this inequality with ν = ρβ and f (a) = c∥a∥2 gives Z Z 2 c ∥a∥2 µ(da) ≤ KL(µ∥ρβ ) + log ec∥a∥ ρβ (a) da. Rd

It remains to check that the Gaussian exponential moment is finite. Since β −1 2 ρβ (a) = Zβ exp − ∥a∥ , 2τ we have

Z e Rd

c∥a∥2

ρβ (a) da = Zβ−1

β 2 exp − − c ∥a∥ da < ∞ 2τ Rd

whenever c ∈ (0, β/(2τ )). Therefore the displayed bound holds, and in particular ∞.

∥a∥2 µ(da) <

Lemma E.2 (Holley–Stroock bounded perturbation) Let µ be a probability measure on Rn satisfying LSI(α), namely, I(ν∥µ) ≥ 2α KL(ν∥µ),

∀ ν ≪ µ.

Let ψ ∈ L∞ (µ) and define dµ̃ := R

eψ dµ, eψ dµ

Osc(ψ) := ess sup ψ − ess inf ψ.

Then µ̃ satisfies LSI(α̃) with α̃ ≥ e− Osc(ψ) α. Equivalently, I(ν∥µ̃) ≥ 2αe− Osc(ψ) KL(ν∥µ̃),

∀ ν ≪ µ̃.

Lemma E.3 (Quantitative Gaussian KL after smoothing) Let ρβ (a) = Zβ−1 exp(−β∥a∥2 /(2τ )) on Rd . Let Y = X + σξ, where σ > 0, ξ ∼ N (0, Id ) is independent of X, and E∥Y ∥2 < ∞. Then ν := Law(Y ) has a density and KL(ν∥ρβ ) ≤

β d E∥Y ∥2 + log Zβ − log(2πeσ 2 ). 2τ 2

(61)

In particular, if σ 2 = 2τ η and E∥Y ∥2 ≤ M , then KL(ν∥ρβ ) ≤

βM d + log Zβ − log(4πeτ η). 2τ 2

(62)

Moreover, if p(a) = Zp−1 eψ(a) ρβ (a) with ∥ψ∥∞ ≤ C, then every probability measure µ satisfying KL(µ∥ρβ ) < ∞ obeys KL(µ∥p) ≤ KL(µ∥ρβ ) + 2C. 50

(63)

Proof [Proof of Lemma E.3] Let 2 −d/2

gσ (z) := (2πσ )

∥z∥2 exp − 2 2σ

be the density of σξ. If µX = Law(X), then the law ν = Law(Y ) has the Lebesgue density Z gσ (y − x) µX (dx). q(y) = Rd

Since gσ > 0, we have q(y) > 0 for every y ∈ Rd . We first justify the entropy lower bound without assuming that X has a density. The joint law of (X, Y ) is PX,Y (dx, dy) = µX (dx) gσ (y − x) dy, whereas PX ⊗ ν is (PX ⊗ ν)(dx, dy) = µX (dx) q(y) dy. Therefore PX,Y ≪ PX ⊗ ν, and the mutual information satisfies 0 ≤ I(X; Y ) := KL(PX,Y ∥PX ⊗ ν) Z Z gσ (y − x) = log gσ (y − x) dy µX (dx) q(y) d d ZR ZR Z = gσ (y − x) log gσ (y − x) dy µX (dx) − Rd

q(y) log q(y) dy

= −h(σξ) + h(Y ). Equivalently, d log(2πeσ 2 ). 2 This argument is understood in the usual extended-real sense; the last display gives a finite lower bound. On the other hand, the Gaussian maximum-entropy inequality under the second-moment constraint gives a finite upper bound on h(Y ), since E∥Y ∥2 < ∞. Hence h(Y ) is a finite real number. Expanding the relative entropy to the Gaussian reference gives Z Z β KL(ν∥ρβ ) = q(y) log q(y) dy − q(y) log ρβ (y) dy = E∥Y ∥2 − h(Y ) + log Zβ . 2τ d d R R h(Y ) ≥ h(σξ) =

Using the entropy lower bound above proves KL(ν∥ρβ ) ≤

β d E∥Y ∥2 + log Zβ − log(2πeσ 2 ), 2τ 2

2 2 which is (61). The bound (62) follows by substituting R σψ = 2τ η and E∥Y ∥ ≤ M . For the bounded-perturbation claim, write Zp = e dρβ . Since ∥ψ∥∞ ≤ C, we have

e−C ≤ Zp ≤ eC , 51

log Zp ≤ C.

Moreover, if KL(µ∥ρβ ) < ∞, then µ ≪ ρβ , and since p is equivalent to ρβ , the following computation is legitimate: Z KL(µ∥p) = KL(µ∥ρβ ) − ψ dµ + log Zp ≤ KL(µ∥ρβ ) + C + C. This proves (63).

Record · ID 224500 · SHA-256 4d4a22a6146e53e9
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.