Wasserstein Proximal Policy Gradient

Zhaoyu Zhu; Shuhan Zhang; Rui Gao; Shuang Li

cs.LGMar 3, 2026

Wasserstein Proximal Policy Gradient

Authors: Zhaoyu Zhu, Shuhan Zhang, Rui Gao, Shuang Li

Organizations: Zhiyuan College, Shanghai Jiao Tong University, Shanghai 200240, China · School of Data Science, The Chinese University of Hong Kong, Shenzhen, Guangdong, China · McCombs School of Business, The University of Texas at Austin, Austin, TX, USA

Abstract

We study policy gradient methods for continuous-action, entropy-regularized reinforcement learning through the lens of Wasserstein geometry. Starting from a Wasserstein proximal update, we derive Wasserstein Proximal Policy Gradient (WPPG) via an operator-splitting scheme that alternates an optimal transport update with a heat step implemented by Gaussian convolution. This formulation avoids evaluating the policy's log density or its gradient, making the method directly applicable to expressive implicit stochastic policies specified as pushforward maps. We establish a global linear convergence rate for WPPG, covering both exact policy evaluation and actor-critic implementations with controlled approximation error. Empirically, WPPG is simple to implement and attains competitive performance on standard continuous-control benchmarks.

Figures & tables

Figure 1: Training curves on MuJoCo continuous control benchmarks: Solid lines denote the mean episodic return, while shaded areas represent the 95% confidence interval computed over 10 independent evaluation runs with different random seeds.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

\varphi_{\sigma}(x)=\frac{1}{(2\pi\sigma^{2})^{d_{a}/2}}\exp\!\left(-\frac{\|x\|^{2}}{2\sigma^{2}}\right)

Appendix

Algorithm 1 Plug-in Mixture Entropy Estimation for Gaussian-Convolved Implicit Policy

Table 3

Table 4

Panel of Figure 4 — Figure 4: Ablation study on $\tau$ (left) and $Latent\ Dimension$ (right).

Figure 6: Additional Ablation on Eta, Action Samples and Latent Dim with HalfCheetah-v5 Task

Figure 7: Additional Ablation on Eta, Action Samples and Latent Dim with Humanoid-v5 Task

Policy Optimization

Policy optimization develops algorithms for improving decision rules in reinforcement learning, language models, and sequential decision systems.

See all 224 papers

Global Convergence of Wasserstein Policy Gradient for Entropy-Regularized Reinforcement Learning
May 25, 2026Zhaoyu Zhu, Rui Gao, Shuang LiOptimization Convergence Analysis Wasserstein Gradient Flows
Shanghai Jiao Tong University · The University of Texas at Austin · The Chinese University of Hong Kong, Shenzhen
A note on convergence of Wasserstein policy optimization
May 21, 2026David Šiška, Yufei ZhangWasserstein Gradient Flows Policy Optimization
School of Mathematics, University of Edinburgh, United Kingdom · Department of Mathematics, Imperial College London, United Kingdom
Wasserstein Policy Gradient for Entropy-Regularized Linear-Quadratic Control
Aug 7, 2026Zhaoyu Zhu, Rui Gao, Shuang LiMaximum Entropy RL Policy Gradient Methods
University of Illinois Urbana-Champaign · The University of Texas at Austin · The Chinese University of Hong Kong, Shenzhen

May 25, 2026J/K move · Enter open · S save

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

Zhaoyu Zhu, Rui Gao, Shuang Li

Shanghai Jiao Tong University · The University of Texas at Austin · The Chinese University of Hong Kong, Shenzhen

Optimization Convergence Analysis Wasserstein Gradient Flows Maximum Entropy RL Policy Gradient Methods Stochastic Optimization Convergence

Open paper AskPDF arXiv

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--Łojasiewicz 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.

Abstract

Policy Optimization

Explore similar work

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

A note on convergence of Wasserstein policy optimization

Wasserstein Policy Gradient for Entropy-Regularized Linear-Quadratic Control