A note on convergence of Wasserstein policy optimization

2026-05-21

AI Summary

This work addresses the lack of theoretical convergence guarantees for Wasserstein Policy Optimization (WPO) in continuous state-action spaces. Within the framework of entropy-regularized Markov decision processes, the authors leverage mean-field theory and tools from information geometry to establish a local log-Sobolev inequality. By analyzing the energy dissipation properties of the value function along Wasserstein gradient flows, they demonstrate that—under the assumption of sufficient regularity of the solution—the value function monotonically decreases and converges linearly to the global optimum. This result provides the first theoretical foundation for the linear convergence of WPO in continuous domains.

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.

Problem

• Wasserstein Policy Optimization
• convergence
• continuous state and action spaces
• entropy-regularised MDPs
• gradient flows

Innovation

• Wasserstein Policy Optimization
• linear convergence
• log-Sobolev inequality
• mean-field analysis
• entropy-regularized MDP