Metric-Oriented Policy Optimization

• Metric-oriented policy optimization is a reinforcement learning framework that prioritizes explicit geometric metrics over traditional divergence measures for defining updates.
• It leverages methodologies such as Wasserstein gradient flows, trust regions with Sinkhorn distances, and Riemannian formulations to design more stable optimization strategies.
• The approach reveals hidden convexity in policy spaces, enabling well-posed algorithms that effectively handle online and offline reinforcement learning challenges.

Searching arXiv for recent and foundational papers on metric-oriented policy optimization, including Wasserstein, trust-region, metric-learning, and Riemannian formulations. Metric-oriented policy optimization denotes a family of reinforcement-learning formulations in which the choice of metric, divergence, or geometric structure is treated as a primary design element of the optimization problem rather than as a secondary implementation detail. Across recent work, this idea appears in several distinct but related forms: policy optimization as a Wasserstein gradient flow on probability measures over actions or parameters (Zhang et al., 2018); trust-region methods that replace Kullback–Leibler divergence with Wasserstein or Sinkhorn distances (Song et al., 2023); policy optimization on a Riemannian manifold induced by the $2$-Wasserstein geometry of action distributions and stationary state distributions (Dus, 16 Apr 2026); parameter-space methods that learn a metric tensor and regularize the divergence of the policy-gradient vector field (Chen et al., 2023); and offline goal-conditioned RL methods that learn a state metric whose distance ordering recovers optimal behavior (Reichlin et al., 2024). A broader implication is that policy optimization can be organized around an explicit notion of distance—between policies, action distributions, occupancies, parameters, or latent states—and that this choice determines gradients, geodesics, Hessians, trust regions, and update stability.

1. Conceptual scope and historical formulations

A central conceptual shift in this literature is to stop treating a policy merely as a parameter vector $\theta$ and instead treat it as a probability measure or a measurable map into a space of probability measures (Zhang et al., 2018). In the indirect formulation of Wasserstein gradient flows, one optimizes a distribution $\mu(\theta)$ over parameters rather than a single parameter vector; in the direct formulation, one optimizes $\pi(\cdot\mid s)$ as a probability distribution over actions for each state (Zhang et al., 2018). In both cases, the search space is a Wasserstein space of probability measures equipped with the $2$-Wasserstein distance $W_2$, and optimization is defined as gradient flow in that metric space (Zhang et al., 2018).

A second line of work formulates metric-oriented policy optimization through trust-region design. Trust-region methods based on Kullback–Leibler divergence are described as pervasive in policy optimization, but the Wasserstein and Sinkhorn extensions replace the usual KL ball with a Wasserstein or entropically regularized optimal-transport ball over action distributions (Song et al., 2023). In this view, the trust region is defined not just in probability space but in a metric space of actions, using a ground cost $d(a,a')$ that measures action similarity (Song et al., 2023).

A third line of work pushes the metric orientation further by endowing the policy space itself with a Riemannian structure. Policies are modeled as measurable maps

$\pi: S \to P_2(A),$

where $P_2(A)$ is the space of probability measures on the action space with finite second moment, equipped with the $2$-Wasserstein distance $\theta$0 (Dus, 16 Apr 2026). The policy manifold is then given a Riemannian metric obtained by averaging Wasserstein tangent-space inner products across states, weighted by the stationary state distribution $\theta$1 induced by the policy (Dus, 16 Apr 2026).

Other papers instantiate the same theme in different spaces. OMPO uses a divergence between transition occupancies as the object that regularizes policy learning under policy or dynamics shifts (Luo et al., 2024). MetricRL uses metric learning over state embeddings so that the learned distance is monotonic with shortest-path distance to goals in deterministic, sparse-reward offline RL (Reichlin et al., 2024). Deep Metric Tensor Regularized Policy Gradient learns a metric tensor $\theta$2 over policy parameters and uses it to smooth the policy-gradient vector field by minimizing its absolute divergence in the induced Riemannian manifold (Chen et al., 2023). Orthogonalized Policy Optimization, in RLHF rather than standard control, explicitly separates sampling geometry from optimization geometry, using an $\theta$3-divergence-inspired sampling weight together with a Bregman-divergence-based value metric (Zixian, 18 Jan 2026). This suggests that “metric-oriented” refers not to a single algorithmic family but to a design principle: choose the geometry first, then define optimization relative to that geometry.

2. Wasserstein geometry on policies and policy distributions

The Wasserstein formulation is the most explicit geometric realization of metric-oriented policy optimization. In “Policy Optimization as Wasserstein Gradient Flows,” the policy optimization problem is placed into the space of probability measures and interpreted as Wasserstein gradient flow (Zhang et al., 2018). For indirect policy learning, the energy functional is

$\theta$4

where $\theta$5; for direct policy learning, at each state $\theta$6,

$\theta$7

with $\theta$8 (Zhang et al., 2018). The associated Wasserstein gradient flow is expressed by the continuity equation

$\theta$9

and the JKO scheme

$\mu(\theta)$0

provides the discrete-time optimization principle (Zhang et al., 2018).

The 2026 Wasserstein formulation generalizes this to a fully Riemannian policy manifold. The action space $\mu(\theta)$1 is a Polish metric space, and

$\mu(\theta)$2

defines the $\mu(\theta)$3-Wasserstein metric on $\mu(\theta)$4 (Dus, 16 Apr 2026). Policies are measurable fields of action distributions,

$\mu(\theta)$5

and the environment induces a stationary measure $\mu(\theta)$6 through the Markov kernel $\mu(\theta)$7 (Dus, 16 Apr 2026). The policy-space Riemannian metric is then

$\mu(\theta)$8

where $\mu(\theta)$9 are tangent velocity fields in action space (Dus, 16 Apr 2026).

This metric is “local” in actions and “global” in states: locally one uses the Wasserstein metric on action distributions; globally one averages over states according to the invariant visitation frequencies $\pi(\cdot\mid s)$0 (Dus, 16 Apr 2026). The framework then defines geodesics, tangent spaces, gradients, and Hessians intrinsically on policy space. In particular, for long-run average cost

$\pi(\cdot\mid s)$1

the Riemannian gradient is

$\pi(\cdot\mid s)$2

where $\pi(\cdot\mid s)$3 is defined through a Poisson/Bellman equation (Dus, 16 Apr 2026). The corresponding Wasserstein gradient flow is

$\pi(\cdot\mid s)$4

so policy improvement becomes transport of probability mass in action space along $\pi(\cdot\mid s)$5 (Dus, 16 Apr 2026).

This directly contrasts with KL/Fisher approaches. The Wasserstein metric is derived from the metric structure of the action space and remains finite and meaningful when supports are disjoint, whereas KL can be infinite or undefined under support mismatch (Dus, 16 Apr 2026, Zhang et al., 2018). A plausible implication is that Wasserstein geometry is especially aligned with continuous control tasks in which Euclidean proximity of actions carries physical meaning.

3. Trust regions defined by Wasserstein, Sinkhorn, and related metrics

Metric-oriented policy optimization also appears as a redesign of trust-region constraints. In Wasserstein Policy Optimization and Sinkhorn Policy Optimization, the local surrogate improvement objective is constrained by a Wasserstein or Sinkhorn distance between successive action distributions rather than by KL divergence (Song et al., 2023). For a discrete action set with ground cost $\pi(\cdot\mid s)$6, the per-state WPO subproblem is

$\pi(\cdot\mid s)$7

The corresponding Lagrangian is

$\pi(\cdot\mid s)$8

and the closed-form update transports the probability mass of each old action $\pi(\cdot\mid s)$9 to the action $2$0 maximizing

$2$1

(Song et al., 2023). Thus each piece of mass is redistributed according to an advantage–distance tradeoff rather than reweighted independently by an exponential function of advantage.

The Sinkhorn version adds entropic regularization to the transport problem. The update becomes soft rather than hard, with transport weights of the form

$2$2

again followed by mass aggregation

$2$3

(Song et al., 2023). The paper proves that WPO guarantees monotonic performance improvement, that SPO converges to WPO as the entropic regularizer vanishes, and that with a decaying Lagrangian multiplier both converge to a globally optimal policy in finite MDPs under exact expectations (Song et al., 2023).

The continuous-time continuous-space RL paper gives an analogous metric-oriented reading of TRPO/PPO in a diffusion setting. It derives performance-difference and local-approximation formulas from a discounted occupation-time measure and extends them to continuous-time PG and TRPO/PPO methods (Zhao et al., 2023). Although the detailed trust-region machinery differs from the discrete setting, the role of the metric is similar: a local approximation of $2$4 is combined with a policy-distance control term.

CIM-PPO modifies PPO at the metric level by replacing KL divergence with the Liu-Correntropy Induced Metric, a symmetric, bounded kernel-induced metric in an RKHS (Guo et al., 2021). The PPO-CIM surrogate is

$2$5

and the motivation is explicit: asymmetry in KL divergence can misguide the PPO-KL surrogate, especially for high-dimensional Gaussian policies with differing variances (Guo et al., 2021). The paper reports better empirical performance than both PPO-KL and PPO-Clip on most of its MuJoCo-style tests (Guo et al., 2021).

A related but distinct unification appears in “A Parametric Class of Approximate Gradient Updates for Policy Optimization,” which organizes many updates around two scalar signals,

$2$6

and a scaling function $2$7 that modulates the update magnitude (Gummadi et al., 2022). PPO corresponds to a specific clipped scale function,

$2$8

while a broader family $2$9 interpolates between value-based, policy-gradient, self-imitation, and PPO-like behaviors (Gummadi et al., 2022). This suggests a metric-oriented reading in which the scaling function defines a sample-wise optimization geometry.

4. Metrics on occupancies, latent states, and parameter space

Not all metric-oriented formulations operate directly on policy distributions. OMPO identifies transition occupancy matching as a unified principle for RL under policy shifts, nonstationary dynamics, and domain adaptation (Luo et al., 2024). The key object is the transition occupancy

$W_2$0

and the surrogate learning objective explicitly penalizes divergence between current and replay occupancies (Luo et al., 2024). After dual reformulation, the policy is trained with a critic and a discriminator estimating a transition-occupancy log-ratio

$W_2$1

which functions as a learned metric on transition distributions (Luo et al., 2024). This is metric-oriented because the policy is not optimized only for reward: it is optimized for reward minus occupancy discrepancy.

MetricRL does something analogous in state space for goal-conditioned offline RL. Under deterministic transitions, invertible actions, sparse rewards, and connectivity assumptions, the optimal value is tied to graph shortest-path distance to the goal (Reichlin et al., 2024). The method learns an embedding

$W_2$2

with Euclidean metric $W_2$3, and introduces the property of distance monotonicity: $W_2$4 If this property holds, greedy control based on latent distance to the goal recovers the optimal policy (Reichlin et al., 2024). The approximate value function is

$W_2$5

and the actor is trained by weighting behavior-cloning updates by $W_2$6, so actions that move closer to the goal in latent metric space are reinforced (Reichlin et al., 2024).

Deep Metric Tensor Regularized Policy Gradient moves the metric orientation to policy parameter space. Standard policy gradient treats $W_2$7 as Euclidean; the paper instead learns a metric tensor $W_2$8 over parameter space and defines the policy-gradient vector field

$W_2$9

then minimizes the squared divergence

$d(a,a')$0

with respect to the metric-network parameters $d(a,a')$1 (Chen et al., 2023). The chosen metric has the structured form

$d(a,a')$2

so the inverse metric yields a preconditioned gradient

$d(a,a')$3

(Chen et al., 2023). The paper then derives a geodesic-regularized gradient and reports that the metric-tensor regularization reduces the absolute divergence of the policy-gradient vector field and improves performance on several benchmarks (Chen et al., 2023).

Orthogonalized Policy Optimization introduces another decomposition: sampling geometry determines which samples dominate the gradient signal, while optimization geometry determines how deviations are penalized (Zixian, 18 Jan 2026). The general objective is expressed as

$d(a,a')$4

and OPO instantiates this with $d(a,a')$5-weighted importance sampling and a Pearson $d(a,a')$6-induced quadratic regularization in ratio coordinates (Zixian, 18 Jan 2026). This makes the metric choice explicit and separates it from the sample-weighting mechanism.

5. Theoretical properties: convexity, Hessians, monotonicity, and well-posedness

One recurring claim in this literature is that choosing the right metric can reveal hidden convexity or better conditioning. In policy optimization as Wasserstein gradient flow, the JKO step

$d(a,a')$7

is convex over measures for the Itô-type energies under the stated assumptions, even though the corresponding parameter optimization problem is highly nonconvex (Zhang et al., 2018). The paper formalizes convergence of the flow to a target distribution induced by return or soft $d(a,a')$8-values (Zhang et al., 2018).

The 2026 Wasserstein manifold paper develops a second-order theory. Along a constant-speed geodesic $d(a,a')$9 with initial velocity $\pi: S \to P_2(A),$0, the Hessian is

$\pi: S \to P_2(A),$1

with the second term capturing Poisson-feedback curvature induced by dynamics (Dus, 16 Apr 2026). In the special case where transitions do not depend on actions, the nonlocal term vanishes and geodesic convexity of $\pi: S \to P_2(A),$2 is equivalent to displacement convexity of the action-cost $\pi: S \to P_2(A),$3 for almost every state (Dus, 16 Apr 2026).

Trust-region papers supply monotonic-improvement guarantees. WPO derives a TRPO-style lower bound and shows that in the exact-advantage setting WPO produces a non-decreasing sequence of returns (Song et al., 2023). SPO converges to WPO as entropic regularization vanishes (Song et al., 2023). CIM-PPO proves that under suitable kernel conditions the new policy remains within a trust region defined by the kernel-induced metric, while also arguing that the metric is computationally cheaper than KL-based policy gradients (Guo et al., 2021).

The metric-tensor paper focuses on divergence and Hessian trace. In Euclidean coordinates, the divergence of the policy-gradient vector field is the trace of the Hessian; by learning a Riemannian metric $\pi: S \to P_2(A),$4 and reducing $\pi: S \to P_2(A),$5, the algorithm attempts to smooth the gradient field in the learned manifold (Chen et al., 2023). This suggests that metric design can be used as a second-order regularizer even when the update remains first order in implementation.

A different notion of well-posedness appears in “Compactly Restrictable Metric Policy Optimization Problems,” which studies deterministic MDPs with metric state and action spaces (Dorobantu et al., 2022). The paper defines a metric policy optimization problem as an MDP on separable metric state and action spaces together with a policy class admitting a partial order, and shows that a compact forward-invariant restriction yields a well-posed Bellman operator, existence of optimal policies, and value-iteration convergence (Dorobantu et al., 2022). In this usage, “metric” refers primarily to the topology and compactness structure required for dynamic programming. This suggests a broader interpretation of metric-oriented policy optimization that includes well-posedness results for continuous-control problems on metric spaces.

6. Practical algorithms, implementations, and broader implications

Several papers translate the geometric ideas into practical algorithms. The Wasserstein-gradient-flow paper uses particle approximations and JKO discretization, giving IP-WGF for parameter distributions and DP-WGF / DP-WGF-V for direct action-distribution updates (Zhang et al., 2018). The Riemannian Wasserstein paper gives both low-dimensional particle methods and high-dimensional deterministic-policy implementations, where the pullback metric on parameters becomes

$\pi: S \to P_2(A),$6

and the parameter update solves

$\pi: S \to P_2(A),$7

via conjugate gradient (Dus, 16 Apr 2026).

MetricRL is practical in offline settings because it never estimates a $\pi: S \to P_2(A),$8-function for out-of-distribution actions; the critic is induced by metric distance, and the actor only reweights observed actions (Reichlin et al., 2024). OMPO remains actor–critic in appearance but uses a distribution discriminator and a small local buffer to estimate transition-occupancy discrepancy under changing policies and dynamics (Luo et al., 2024). The metric-tensor method augments SAC and TD3 with an inner loop that trains the metric tensor network before each policy update (Chen et al., 2023). CIM-PPO keeps the overall PPO structure and simply replaces the KL penalty by a kernel metric, using triangular or biweight kernels in the reported experiments (Guo et al., 2021). The parametric family of approximate gradient updates can be dropped into PPO-style code by replacing the standard clipped scale with $\pi: S \to P_2(A),$9 (Gummadi et al., 2022).

The literature also contains metric-oriented formulations outside standard control RL. MePoly is framed around maximum-entropy stochastic control and explicit information-theoretic metrics such as Shannon entropy and KL divergence, using a polynomial energy-based model to keep $P_2(A)$0, entropy, and KL tractable for highly multimodal policies (Liu et al., 19 Feb 2026). Orthogonalized Policy Optimization treats the divergence or Bregman geometry itself as a first-class design variable in RLHF (Zixian, 18 Jan 2026). A plausible implication is that metric-oriented policy optimization is becoming a general design language for both reinforcement learning and alignment.

There are also objective cautions. The paper on Metric Aggregation Divergence argues that in agent-based model plus multi-objective evolutionary pipelines, different stages can silently use inconsistent implementations of the same metric, leading to different “champion” policies and different statistical conclusions (Zhang et al., 27 Jun 2026). Its proposed remedy is a metric contract: a single shared callable enforced across optimization, post-processing, and inference (Zhang et al., 27 Jun 2026). Although this work is outside standard policy-gradient RL, it highlights that metric-oriented optimization is only as reliable as the consistency of the metric across the pipeline.

Overall, the research record supports a common conclusion. Metric-oriented policy optimization is not simply “policy optimization with a different regularizer.” It is a perspective in which the geometry of policy space, action space, occupancy space, latent state space, or parameter space is made explicit and then used to define the optimization problem itself. Depending on the chosen metric, this yields Wasserstein gradient flows (Zhang et al., 2018, Dus, 16 Apr 2026), Wasserstein/Sinkhorn trust regions (Song et al., 2023), kernel-metric PPO variants (Guo et al., 2021), occupancy-matching objectives (Luo et al., 2024), latent-metric offline RL (Reichlin et al., 2024), learned Riemannian parameter manifolds (Chen et al., 2023), or explicit decompositions of sampling and optimization geometry (Zixian, 18 Jan 2026). This suggests that future work will continue to treat metric choice not as a technical detail, but as one of the central modeling decisions in policy optimization.