A note on convergence of Wasserstein policy optimization

Abstract: Wasserstein Policy Optimization (WPO) is a recently proposed reinforcement learning algorithm that leverages Wasserstein gradient flows to optimize stochastic policies in continuous action spaces. Despite its empirical success, the theoretical convergence properties of WPO in environments with continuous state and action spaces have yet to be fully established. In this note, we argue that WPO within the framework of entropy-regularised Markov Decision Processes converges linearly. This is done by leveraging recent advances in mean-field analysis for convergence of gradient flows using log-Sobole inequalities. Assuming existence of sufficiently regular solution to the gradient flow equation we demonstrate monotonic energy dissipation along the flow and establish a local log-Sobolev inequality. Ultimately, these properties allow us to argue that the value function should converge linearly to the global optimum.

• The paper presents a rigorous convergence analysis demonstrating global linear (exponential) convergence of the value function in entropy‐regularised MDPs using Wasserstein gradient flows.
• It leverages log-Sobolev inequalities and energy dissipation techniques to quantify policy improvement and ensure continuous, non-stagnating progress.
• The analysis offers strong theoretical foundations and practical insights for extending WPO to discretized algorithms and high-dimensional continuous-action settings.

Convergence Analysis of Wasserstein Policy Optimization in Entropy-Regularised MDPs

Introduction

This paper rigorously investigates the convergence behavior of Wasserstein Policy Optimization (WPO) within the framework of entropy-regularised Markov Decision Processes (MDPs) featuring infinite-horizon, discrete state, and continuous action spaces. WPO employs elements of optimal transport—specifically, Wasserstein gradient flows—to optimize stochastic policy distributions in reinforcement learning (RL). While WPO demonstrates strong empirical results in continuous control domains, prior theoretical work has not resolved its convergence guarantees. This study addresses that gap by providing a linear convergence rate analysis for WPO, leveraging tools from convex mean-field dynamics, log-Sobolev inequalities, and the analysis of energy dissipation in gradient flows.

Problem Formulation and Mathematical Setting

The setting is an infinite-horizon entropy-regularised MDP $(S, A, P, c, \gamma)$, with $S$ (discrete), $A=\mathbb{R}^d$ (continuous), a bounded cost $c$, transition kernel $P$, and discount factor $\gamma\in[0,1)$. Policies are stochastic and potentially history-dependent. Entropy regularisation is introduced via a Kullback-Leibler (KL) term with respect to a reference measure $p$. The regularised value function $V^\pi(s)$ combines long-term costs and an entropy term, and the optimal value function $V^*$ as well as the associated Bellman equations are given explicit forms.

The policy space is parameterized as an exponential family with respect to the reference measure, and the analysis focuses on gradient flows of the value functional with respect to the Wasserstein metric on probability measures over actions. The WPO procedure is viewed as a continuous-time process described by a continuity equation, with the policy distribution evolving under Wasserstein gradient flow driven by the flat derivative (functional gradient) of the regularised value function.

Key Theoretical Contributions

Energy Dissipation and Log-Sobolev Inequality

A central technical component is establishing monotonic energy (value function) dissipation along the gradient flow trajectory under mild regularity assumptions. The work builds toward a local log-Sobolev inequality that holds along the flow, differing from prior literature by not requiring uniformity over the entire policy space. The log-Sobolev inequality is instrumental in quantifying the exponential rate at which the relative entropy between successive policies decays, directly relating to the speed of value function improvement.

Linear (Exponential) Convergence Rate

The principal result asserts that the value function $V^\pi$ converges linearly (that is, exponentially in continuous time) to the global optimum $S$0, provided the solution to the continuity equation governing the Wasserstein gradient flow exists with the required regularity:

$S$1

where $S$2 and $S$3 are constants dependent on the occupancy measure and log-Sobolev constants, and $S$4 is the occupancy measure for the optimal policy.

The derivation rests on demonstrating:

• An “entropy sandwich” inequality that relates the difference in value functions to KL divergences between policy distributions.
• Strict energy dissipation via Wasserstein gradient flow.
• Existence of local log-Sobolev inequalities along the flow, enabling the application of Gronwall-type arguments for exponential contraction.
• Suitable occupancy measure bounds for the convergence constants.

These results collectively guarantee that the optimization trajectory never stagnates and always makes analytic progress toward the global optimum.

Policy Space and Measure-Theoretic Tools

The formalization relies on advanced measure-theoretic constructs, including the duality pairing for value differences, occupancy kernel analysis, and careful handling of non-differentiability and distribution shifts arising during policy evolution.

Numerical and Qualitative Implications

Although the results are proved in continuous time, the authors state that analogous guarantees should hold for discretised/iterative versions of the algorithm, contingent on extending arguments from recent work on discrete Wasserstein gradient schemes. The convergence is global and linear with respect to the value function, ensuring effective optimization even as dimensionality and policy class richness increase—subject to regularity and log-Sobolev preconditions.

The work makes the strong claim that, in contrast to most policy gradient or actor-critic algorithms whose general convergence rates are sublinear and often local, WPO achieves provable global linear convergence under entropy-regularisation. This claim is substantiated by reducing the problem to established results in convex mean-field analysis but adapting them to the specific technical challenges of the RL setting.

Theoretical and Practical Implications

From a theoretical perspective, this analysis strengthens the mathematical foundations for using Wasserstein-based optimization methods in RL. The techniques developed, particularly those regarding the interplay between gradient flow, entropy, and log-Sobolev inequalities, clarify the conditions under which fast convergence can be expected in high-dimensional, continuous-action MDPs.

Practically, this result implies that WPO is not only empirically effective but also robust from an optimization perspective, provided entropy regularisation is used and regularity assumptions are met. Future developments may include extending these convergence proofs to:

• Practical discrete-time (step-wise) WPO algorithms, which are necessary for actual computation,
• More general state spaces (beyond discrete),
• Relaxed regularity conditions, possibly via functional approximation methods such as neural policies,
• Alternative regularisers beyond entropy.

Furthermore, this line of work could motivate RL algorithm designers to prioritise regularisation and optimal transport-based methods for scalable, reliable policy optimization in continuous domains.

Conclusion

This study provides a rigorous analysis demonstrating that Wasserstein Policy Optimization enjoys global linear convergence in entropy-regularised MDPs under appropriate assumptions. The results bridge mean-field optimal transport theory and RL, yielding strong guarantees for policy improvement when WPO is applied with continuous actions and entropy regularization. This positions WPO as a theoretically well-grounded method and paves the way for further mathematically robust policy optimization techniques in high-dimensional reinforcement learning settings.