Optimal Transport Geodesic Picard Flow

• OTGP flow is a continuous-time mechanism that evolves probability measures along Wasserstein-2 geodesics, ensuring smooth transitions in mean-field updates.
• It integrates optimal transport theory with Picard iteration to update distributions within actor-critic frameworks, enhancing stability and convergence.
• The flow achieves global exponential convergence, validated by Lyapunov analysis and demonstrated through applications in high-dimensional equilibrium problems.

Optimal Transport Geodesic Picard (OTGP) Flow refers to a continuous-time flow in the space of probability measures that updates a candidate distribution toward a target distribution by moving along Wasserstein-2 geodesics, leveraging principles from optimal transport theory and fixed-point (Picard) iteration. The OTGP flow concept arises in the context of mean-field reinforcement learning algorithms and mean-field games, most prominently as a component of the Mean-Field Actor-Critic (MFAC) framework where it is used as the distribution update mechanism that ensures convergence of the empirical distribution toward equilibrium (Zhou et al., 14 Oct 2025).

1. Concept and Role in Learning Dynamics

The central objective of the OTGP flow is to realize a continuous update from a current measure (distribution) toward a target (typically the distribution induced by the current control in a mean-field game). Unlike classical discrete “fictitious play” updates that replace the candidate measure in each iteration, the OTGP flow interpolates between the measures via a geodesic in the Wasserstein metric space. This produces a smooth path of distributions that respects the geometric structure of the space of probability measures and ensures consistency with the state population generated by the current control policy.

Within coupled learning frameworks such as MFAC, the OTGP flow serves as one of three tightly coupled PDE-based update mechanisms:

• The actor (policy update) via policy gradient flow,
• The critic (value function estimation) via a shooting method,
• The distribution component, i.e., the empirical measure, updated by the OTGP flow.

The OTGP flow ensures that the current distribution μ(τ) evolves toward the distribution ρ induced by the current policy, along the Wasserstein-2 geodesic.

2. Mathematical Formulation

The OTGP flow is characterized by its evolution along Wasserstein-2 geodesics. For any fixed time t and learning-time parameter τ, let μᵗ(τ) be the current measure and ρₜ the target (as generated by the actor and critic). The optimal transport map Tᵗ(τ) from μᵗ(τ) to ρₜ is constructed using the Kantorovich potential ϕᵗ(τ, x):

$Tᵗ(τ, x) = x - ∇ϕᵗ(τ, x)$

The time evolution of the candidate distribution μᵗ(τ) is governed by the continuity equation

$$∂_τ μᵗ(τ,x) + ∇_x \cdot \left( μᵗ(τ,x) [-βₘ ∇ϕᵗ(τ,x)] \right) = 0$$

where βₘ is a positive update rate. Equivalently, in discrete time with step size Δτ,

$$μ_{τ+Δτ}ₜ \approx [\mathrm{Id} - Δτ · βₘ ∇ϕᵗ(τ,·)]_\# μᵗ(τ)$$

This update moves the distribution μᵗ(τ) toward ρₜ along the Wasserstein-2 geodesic. The flow acts as a Picard iteration in the Wasserstein metric: repeatedly applying the map $\mathrm{Id} - Δτ · βₘ ∇ϕ$ to the measure drives μ toward the target ρ.

3. Convergence Analysis

Global exponential convergence of the OTGP flow is established via Lyapunov functionals. Define the function

$$Lₘ(τ) = \int_0^T e^{-2β t} W_2\big(\muₜ(τ),\,ρₜ\big)^2\,dt + ½ λ_T W_2(\mu_T(τ),\,ρ_T)^2$$

where W₂ denotes the squared Wasserstein-2 distance, and λ_T, β are positive weights. The OTGP flow, when incorporated into coupled learning where the actor and critic errors are controlled, produces a negative-definite derivative:

$$\frac{d}{dτ}Lₘ(τ) \le -cₘ βₘ Lₘ(τ) + (\text{coupling error terms})$$

for some positive constant cₘ. By properly choosing update rates so that the distribution step evolves much faster than the actor and critic, one guarantees that Lₘ(τ) converges exponentially fast to zero, implying convergence of μ(τ) to ρ in Wasserstein metric as τ → ∞.

4. Algorithmic Coupling and System Interplay

In MFAC, the OTGP flow is fully coupled with both actor and critic updates:

• The actor relies on the current empirical distribution to evaluate its policy gradient.
• The critic depends on the current candidate measure for value function estimation.
• The distribution update (OTGP flow) moves the measure μ toward the distribution generated by the policy (ρ) via a sequence of geodesic interpolations.

The mutual dependence is resolved by joint evolution: the actor and critic are updated with their own learning dynamics (policy gradient and PDE-based shooting), while the OTGP flow keeps the current empirical distribution close to the equilibrium induced by the actor’s policy.

5. Discrete Implementation and Computational Steps

The practical implementation of the OTGP flow in MFAC is sample-based:

• Draw particles from the current candidate distribution μᵗ(τ), denoted L^{k,m,t}.
• Compute the optimal transport map Tᵗ(τ,·) (using the Kantorovich dual, solved via the Hungarian algorithm).
• Update the sample ensemble as

$$Q^{k+1,m}_t = Δτ·βₘ T^\tau(L^{k,m,t}) + (1-Δτ·βₘ) L^{k,m,t}$$

• Parallelly update the actor and critic blocks.
• Empirically observed, each distributional update requires only a small number of iterations to contract to equilibrium.

6. Numerical Results and Effectiveness

Empirical evaluations on mean-field equilibrium problems—such as systemic risk, optimal execution in high dimensions, and high-dimensional flocking dynamics—demonstrate that the OTGP flow component successfully drives the candidate distribution to the MFG equilibrium. In these examples, the Wasserstein-2 distance between the current empirical distribution and the analytic/numerical ground-truth solution rapidly decays, and the sample-based estimated densities agree with benchmark solutions.

7. Connections and Significance

The OTGP flow generalizes fixed-point (Picard) iteration schemes for probability measures, embedding them in the geometric context of optimal transport. Unlike classical approaches that may replace distributions in discrete jumps, the OTGP flow interpolates via geodesics, preserving deep geometric and stability properties. Its Lyapunov structure and contractivity in weighted Wasserstein metrics provide strong convergence guarantees, and it flexibly accommodates high-dimensional and coupled learning dynamics. The continuous, geodesic approach contrasts with both classical fictitious play in mean-field games and “external” marginal relaxation techniques in computational optimal transport: OTGP flow’s monotonic, “interior” movement in measure space ensures robust, interpretable, and globally convergent learning dynamics in high-dimensional, coupled reinforcement learning systems (Zhou et al., 14 Oct 2025).