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

📄 arXiv: 2604.14765v1 📥 PDF

作者: Mathias Dus

分类: cs.LG, math.OC, math.PR

发布日期: 2026-04-16

💡 一句话要点

提出基于Wasserstein空间的强化学习框架，优化策略。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 强化学习 Wasserstein空间 最优传输 策略优化 黎曼几何

📋 核心要点

  1. 现有强化学习方法在处理复杂策略空间时面临挑战，难以保证优化过程的稳定性和效率。
  2. 论文将策略视为Wasserstein空间中的映射，利用最优传输理论的几何结构来优化策略。
  3. 通过数值实验验证了该方法在低维问题上的有效性，并探索了其在高维问题中的应用潜力。

📝 摘要（中文）

本文提出了一个强化学习（RL）的几何框架，将策略视为映射到动作概率Wasserstein空间的映射。首先，我们定义了一个由平稳分布引起的黎曼结构，并证明了它在一般上下文中的存在性。然后，我们定义了策略的切空间，并描述了测地线，特别关注了从状态空间映射到动作空间概率测度切向量场的向量场的可测量性。接下来，我们提出了一个通用的RL优化问题，并使用Otto微积分构建了一个梯度流。我们计算了能量的梯度和Hessian矩阵，提供了一个正式的二阶分析。最后，我们用低维问题的数值例子说明了该方法，直接从我们的理论形式计算梯度。对于高维问题，我们使用神经网络参数化策略，并基于成本的遍历近似来优化它。

🔬 方法详解

问题定义：强化学习旨在寻找最优策略，以最大化累积奖励。现有方法，如策略梯度和Q学习，在处理连续动作空间或高维状态空间时，面临着样本效率低、易陷入局部最优等问题。此外，对策略空间的几何结构缺乏有效利用，限制了优化算法的性能。

核心思路：论文的核心在于将策略空间视为动作概率的Wasserstein空间，并利用最优传输理论的几何结构来定义策略的梯度和Hessian。Wasserstein空间提供了一种度量概率分布之间距离的自然方式，能够更好地捕捉策略之间的相似性。通过在Wasserstein空间中进行策略优化，可以更有效地探索策略空间，并提高算法的稳定性和收敛速度。

技术框架：该方法首先定义了由平稳分布诱导的黎曼结构，并证明了其存在性。然后，定义了策略的切空间，并刻画了测地线。接着，构建了一个通用的强化学习优化问题，并利用Otto微积分构建梯度流。最后，计算了能量的梯度和Hessian，提供了一个正式的二阶分析。对于高维问题，采用神经网络参数化策略，并基于成本的遍历近似进行优化。

关键创新：该方法最重要的创新点在于将强化学习问题置于Wasserstein空间中进行分析和优化。与传统的欧几里得空间相比，Wasserstein空间能够更好地捕捉策略之间的相似性，并提供更丰富的几何结构。此外，利用Otto微积分构建梯度流，能够更有效地进行策略优化。

关键设计：论文的关键设计包括：1) 定义了由平稳分布诱导的黎曼结构；2) 利用Otto微积分构建梯度流；3) 对于高维问题，采用神经网络参数化策略，并基于成本的遍历近似进行优化。具体参数设置和网络结构的选择取决于具体的应用场景。

📊 实验亮点

论文通过数值实验验证了该方法在低维问题上的有效性，并探索了其在高维问题中的应用潜力。实验结果表明，该方法能够有效地计算策略的梯度，并提高算法的收敛速度。在高维问题中，通过神经网络参数化策略，并基于成本的遍历近似进行优化，取得了初步的成果。

🎯 应用场景

该研究成果可应用于机器人控制、自动驾驶、金融交易等领域。通过利用Wasserstein空间的几何结构，可以更有效地优化策略，提高智能体的性能和鲁棒性。未来，该方法有望在更复杂的强化学习任务中发挥重要作用，例如多智能体协作、元学习等。

📄 摘要（原文）

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.