Wasserstein Formulation of Reinforcement Learning. An Optimal Transport Perspective on Policy Optimization

• URL: http://arxiv.org/abs/2604.14765v1
• Date: Thu, 16 Apr 2026 08:24:23 GMT
• Title: Wasserstein Formulation of Reinforcement Learning. An Optimal Transport Perspective on Policy Optimization
• Abstract summary: We present a framework for Reinforcement Learning that views policies as maps into the Wasserstein space of action probabilities.<n>For high-dimensional problems, we parameterize the policy using a neural network and optimize it based on an ergodic approximation of the cost.
• Score: 0.0
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: We present a geometric framework for Reinforcement Learning (RL) that views policies as maps into the Wasserstein space of action probabilities. First, we define a Riemannian structure induced by stationary distributions, proving its existence in a general context. We then define the tangent space of policies and characterize the geodesics, specifically addressing the measurability of vector fields mapped from the state space to the tangent space of probability measures over the action space. Next, we formulate a general RL optimization problem and construct a gradient flow using Otto's calculus. We compute the gradient and the Hessian of the energy, providing a formal second-order analysis. Finally, we illustrate the method with numerical examples for low-dimensional problems, computing the gradient directly from our theoretical formalism. For high-dimensional problems, we parameterize the policy using a neural network and optimize it based on an ergodic approximation of the cost.

Related papers

• Fitting Unknown Number of Hyperplanes with Manifold Optimization [57.48093263119306]
  Fitting unknown number of linearplanes to data is a fundamental yet challenging problem in machine learning.<n>Existing approaches often struggle with optimum optima or lack geometric consistency.
  arXiv  Detail & Related papers  (2026-05-27T14:02:20Z)
• Static and Dynamic Approaches to Computing Barycenters of Probability Measures on Graphs [3.702642055407484]
  The optimal transport problem defines a geometry of probability measures which leads to a definition for weighted averages (barycenters) of measures.<n>Here, we implement a barycentric coding model for measures supported on a graph, in which the classical optimal transport geometry becomes degenerate.<n>We conclude that intrinsic gradient descent on the probability simplex provides a coherent framework for the synthesis and analysis of measures supported on graphs.
  arXiv  Detail & Related papers  (2026-03-27T19:29:28Z)
• Riemannian gradient descent for Hartree-Fock theory [0.0]
  We present an optimization framework for Hartree-Fock theory formulated directly in the Sobolev space $H1$.<n>Explicit expressions for Euclidean and Riemannian gradients, tangent-space projections, and retractions are derived.<n>The proposed formulation provides a geometrically consistent and discretization-independent perspective on electronic structure optimization.
  arXiv  Detail & Related papers  (2026-03-16T19:58:42Z)
• Neural Local Wasserstein Regression [16.52489456261937]
  We study the estimation problem of distribution-on-distribution regression, where both predictors and responses are probability measures.<n>Existing approaches typically rely on a global optimal transport map or tangent-space linearization.<n>We propose a flexible nonparametric framework that models regression through locally defined transport maps in Wasserstein space.
  arXiv  Detail & Related papers  (2025-11-13T21:54:18Z)
• Learning Optimal Deterministic Policies with Stochastic Policy Gradients [62.81324245896716]
  Policy gradient (PG) methods are successful approaches to deal with continuous reinforcement learning (RL) problems. In common practice, convergence (hyper)policies are learned only to deploy their deterministic version. We show how to tune the exploration level used for learning to optimize the trade-off between the sample complexity and the performance of the deployed deterministic policy.
  arXiv  Detail & Related papers  (2024-05-03T16:45:15Z)
• Warped geometric information on the optimisation of Euclidean functions [43.43598316339732]
  We consider optimisation of a real-valued function defined in a potentially high-dimensional Euclidean space. We find the function's optimum along a manifold with a warped metric. Our proposed algorithm, using 3rd-order approximation of geodesics, tends to outperform standard Euclidean gradient-based counterparts.
  arXiv  Detail & Related papers  (2023-08-16T12:08:50Z)
• Last-Iterate Convergence of Adaptive Riemannian Gradient Descent for Equilibrium Computation [52.73824786627612]
  This paper establishes new convergence results for textitgeodesic strongly monotone games.<n>Our key result shows that RGD attains last-iterate linear convergence in a textitgeometry-agnostic fashion.<n>Overall, this paper presents the first geometry-agnostic last-iterate convergence analysis for games beyond the Euclidean settings.
  arXiv  Detail & Related papers  (2023-06-29T01:20:44Z)
• Deep Metric Tensor Regularized Policy Gradient [4.982806898121435]
  Policy gradient algorithms are an important family of deep reinforcement learning techniques. We conduct research driven by the believe that properly utilizing and controlling Hessian information associated with the policy gradient can noticeably improve the performance of policy gradient algorithms.
  arXiv  Detail & Related papers  (2023-05-18T14:50:00Z)
• Learning High Dimensional Wasserstein Geodesics [55.086626708837635]
  We propose a new formulation and learning strategy for computing the Wasserstein geodesic between two probability distributions in high dimensions. By applying the method of Lagrange multipliers to the dynamic formulation of the optimal transport (OT) problem, we derive a minimax problem whose saddle point is the Wasserstein geodesic. We then parametrize the functions by deep neural networks and design a sample based bidirectional learning algorithm for training.
  arXiv  Detail & Related papers  (2021-02-05T04:25:28Z)
• Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization [94.18714844247766]
  Wasserstein barycenters provide a geometric notion of the weighted average of probability measures based on optimal transport. We present a scalable algorithm to compute Wasserstein-2 barycenters given sample access to the input measures.
  arXiv  Detail & Related papers  (2021-02-02T21:01:13Z)
• On Projection Robust Optimal Transport: Sample Complexity and Model Misspecification [101.0377583883137]
  Projection robust (PR) OT seeks to maximize the OT cost between two measures by choosing a $k$-dimensional subspace onto which they can be projected. Our first contribution is to establish several fundamental statistical properties of PR Wasserstein distances. Next, we propose the integral PR Wasserstein (IPRW) distance as an alternative to the PRW distance, by averaging rather than optimizing on subspaces.
  arXiv  Detail & Related papers  (2020-06-22T14:35:33Z)
• Distributed Averaging Methods for Randomized Second Order Optimization [54.51566432934556]
  We consider distributed optimization problems where forming the Hessian is computationally challenging and communication is a bottleneck. We develop unbiased parameter averaging methods for randomized second order optimization that employ sampling and sketching of the Hessian. We also extend the framework of second order averaging methods to introduce an unbiased distributed optimization framework for heterogeneous computing systems.
  arXiv  Detail & Related papers  (2020-02-16T09:01:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.