Policy Mirror Descent Overview

• Policy Mirror Descent is a policy optimization method that performs statewise convex mirror steps using evaluation signals to update policies.
• It generalizes exponentiated gradient updates, recovering traditional policy iteration in the limit while ensuring linear convergence under proper step sizes.
• PMD’s adaptable mirror maps, lookahead variants, and robust sample-complexity theory make it applicable to diverse settings including multi-agent and continuous control.

Policy Mirror Descent (PMD) is a policy-optimization framework that performs proximal, or “mirror,” steps in policy space using policy-evaluation information. In discounted Markov decision processes, it updates a policy by solving a statewise convex optimization problem over the action simplex, where a Bregman divergence induced by a mirror map controls the geometry of the update. In its unregularized form, PMD optimizes the original discounted return and uses the proximal term only to regularize the policy-improvement step; with the negative-entropy mirror map it yields the exponentiated-gradient update commonly associated with natural policy gradient, and in the limit of infinite step size it recovers policy iteration (Johnson et al., 2023, Protopapas et al., 2024).

1. Formalism in policy space

A standard discounted finite MDP consists of a finite state space $S$, a finite action space $A$, a transition kernel $P$, a reward $r(s,a)\in[0,1]$, and a discount factor $\gamma\in(0,1)$. For a stationary policy $\pi$, the value function and action-value function are

$$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$

The advantage is $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$, and the discounted visitation distribution is

$$d^\pi_{\rho}(s)=(1-\gamma)\sum_{t=0}^\infty \gamma^t \mathbb{P}^\pi(s_t=s\mid s_0\sim \rho).$$

A central identity is the tabular policy gradient

$$\frac{\partial}{\partial \pi(a|s)}V^\pi(\rho)=\frac{1}{1-\gamma}d^\pi_\rho(s)Q^\pi(s,a),$$

which motivates state-weighted mirror steps (Johnson et al., 2023).

With a Legendre mirror map $A$0 and associated Bregman divergence $A$1, PMD updates each state distribution by

$A$2

This is the policy-space proximal form. In the full-policy view, PMD uses a weighted Bregman divergence $A$3, and the visitation weights cancel against the policy gradient, leaving independent statewise subproblems (Johnson et al., 2023).

Different mirror maps induce different geometries. With the negative-entropy mirror map, $A$4 is the Kullback–Leibler divergence and the update has the closed form

$A$5

With Euclidean geometry, the update becomes a projection onto the simplex; with Tsallis-type generators one obtains generalized entropy geometries, although the closed form may be intractable for general $A$6 (Johnson et al., 2023, Lin et al., 2022).

A broader parameterized view appears in Approximate Mirror Policy Optimization (AMPO), where policies are defined by a Bregman projection of a score function: $A$7 For the negative-entropy mirror map this recovers the softmax class, including tabular softmax, log-linear policies, and neural softmax policies; for $A$8 geometry it yields Euclidean projection onto the simplex (Alfano et al., 2023).

2. Relation to policy iteration, trust regions, and trajectory-space methods

PMD is often described as soft policy iteration. Classical policy iteration alternates exact evaluation of $A$9 and greedy improvement $P$0. PMD replaces the hard greedy step by a mirror step that interpolates between the current policy and the greedy policy. This algorithmic regularization is motivated by the instability of policy iteration under inexact evaluation: the proximal term controls the move size and the geometry of policy change while preserving the unregularized control objective (Johnson et al., 2023).

In the exact-evaluation setting, PMD and policy iteration are closely linked. Policy iteration converges linearly with rate $P$1 in value metrics and reaches an optimal policy in finitely many iterations in discounted tabular MDPs. PMD recovers policy iteration in the limit $P$2, but for finite $P$3 it remains a genuine mirror-descent method in policy space (Johnson et al., 2023).

The trust-region interpretation is explicit in the KL geometry. A canonical PMD update may be written as

$P$4

This is equivalent to a KL-penalized trust-region step. In this sense, TRPO is a constrained trust-region form of the same geometry, PPO is a practical relaxation of KL-regularized mirror descent, and natural policy gradient corresponds to the KL or Fisher geometry induced by negative entropy (Protopapas et al., 2024).

A related trajectory-space interpretation appears in guided policy search. There, a teacher trajectory optimizer performs a KL-constrained local policy update, and a student policy is obtained by projection onto the parameterized policy manifold. This yields an approximate mirror-descent view in trajectory space, with exact mirror-descent guarantees in linear-quadratic settings and bounded projection error in nonlinear settings (Montgomery et al., 2016).

3. Exact convergence, optimality, and policy identification

A central exact-analysis result shows that unregularized PMD attains the dimension-free $P$5-rate of policy iteration in discounted finite MDPs. If the step sizes satisfy

$P$6

then

$P$7

With $P$8, this becomes

$P$9

and for rewards in $r(s,a)\in[0,1]$0 with $r(s,a)\in[0,1]$1,

$r(s,a)\in[0,1]$2

The bound is dimension-free in the sense that neither the rate nor the constants depend on $r(s,a)\in[0,1]$3 or $r(s,a)\in[0,1]$4 (Johnson et al., 2023).

The same work establishes non-asymptotic optimality and step-size necessity. For every $r(s,a)\in[0,1]$5, there exist MDPs with $r(s,a)\in[0,1]$6 such that for any choice of PMD step sizes and all $r(s,a)\in[0,1]$7,

$r(s,a)\in[0,1]$8

Thus the $r(s,a)\in[0,1]$9-rate is unimprovable for PMD in the early-iteration regime. For natural policy gradient with negative entropy and $\gamma\in(0,1)$0, a matching necessity theorem shows that achieving the $\gamma\in(0,1)$1-rate requires adaptive step sizes depending on the KL divergence to a greedy distribution; fixed non-adaptive rules can fail on some MDPs (Johnson et al., 2023).

An important analytical feature is that this convergence proof avoids the performance difference lemma. Instead it uses a three-point descent inequality and the direct bound

$\gamma\in(0,1)$2

which yields a sup-norm recursion without distribution-mismatch coefficients (Johnson et al., 2023).

Other theoretical lines extend this picture. In regularized RL, PMD achieves global linear convergence for strongly convex regularizers, and Approximate PMD or Generalized PMD obtain global linear convergence for general convex regularizers by using adaptive perturbation or a Bregman divergence matched to the regularizer (Lan, 2021, Zhan et al., 2021). Homotopic PMD adds a diminishing regularization term and establishes global linear convergence, local superlinear convergence, and last-iterate convergence to the unique optimal policy with maximal entropy; for some decomposable Bregman divergences, including Euclidean-type geometries, finite-time exact convergence is obtained (Li et al., 2022, Lin et al., 2022).

4. Sampling, approximation, and general policy classes

PMD has a parallel sample-complexity theory. Under a generative model, inexact PMD estimates $\gamma\in(0,1)$3 by Monte Carlo rollouts of horizon $\gamma\in(0,1)$4 and uses the same mirror update with $\gamma\in(0,1)$5. With adaptive step sizes based on the estimated greedy set, one obtains, with probability at least $\gamma\in(0,1)$6,

$\gamma\in(0,1)$7

and total sample complexity

$\gamma\in(0,1)$8

This bound is dimension-optimal in the sense that it removes dependence on distribution-mismatch coefficients such as $\gamma\in(0,1)$9 (Johnson et al., 2023).

Under online Markovian sampling with temporal-difference critics, Expected TD-PMD and Approximate TD-PMD replace generative access by single-trajectory data. With a sufficiently small constant policy-update step size, they attain $\pi$0 sample complexity for average-time $\pi$1-optimality, and adaptive policy step sizes improve this to $\pi$2 for last-iterate $\pi$3-optimality (Li et al., 30 Dec 2025). A related on-policy theory shows that stochastic PMD can achieve $\pi$4 without explicit action-space exploration. Value-based estimation is tailored to KL geometry, while truncated on-policy Monte Carlo combined with suitable Bregman divergences, including Tsallis divergences, yields high-probability optimality-gap control and an effective-horizon term $\pi$5 that is only polynomial in the action-space size for appropriate divergences (Li et al., 2023).

PMD is not confined to tabular closure. AMPO decomposes each iteration into a regression step

$\pi$6

followed by a Bregman projection

$\pi$7

Under approximation and concentrability assumptions, this yields the first linear-convergence guarantee for a policy-gradient-based method with general parameterization, and with two-layer ReLU networks the sample complexity is

$\pi$8

up to a non-vanishing approximation floor (Alfano et al., 2023).

A complementary non-tabular theory replaces compatible function approximation and closure assumptions by a strictly weaker variational gradient dominance condition. In that framework, PMD is analyzed as smooth non-convex optimization in a local norm induced by the current occupancy measure, yielding best-in-class sublinear convergence over convex policy classes (Sherman et al., 16 Feb 2025). This suggests that PMD’s geometry can be studied independently of the tabular closure conditions that underlie many earlier analyses.

5. Variants, block updates, lookahead, and learned geometry

Several PMD variants modify the policy-improvement operator itself. The most direct is $\pi$9-PMD, which replaces one-step greedy improvement by an $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$0-step lookahead policy. Its update uses $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$1 instead of $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$2, and exact $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$3-PMD satisfies the dimension-free bound

$$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$4

Thus the contraction factor becomes $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$5, and the same paper extends the theory to Monte Carlo lookahead and linear function approximation, where the error depends on feature dimension rather than $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$6 (Protopapas et al., 2024).

Block Policy Mirror Descent replaces full-state updates by a single sampled-state update. Under KL geometry, the update only changes $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$7 at the sampled state $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$8. In the deterministic tabular setting this reduces per-iteration policy-evaluation cost from $$V^\pi(s)=\mathbb{E}\Big[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\mid s_0=s,\pi\Big], \qquad Q^\pi(s,a)=r(s,a)+\gamma\sum_{s'}p(s'|s,a)V^\pi(s').$$9 to $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$0, and under suitable exploratory sampling it achieves linear convergence for strongly convex regularizers. In the stochastic generative-model setting, sample complexity becomes $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$1 for strongly convex regularizers and $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$2 for non-strongly-convex regularizers (Lan et al., 2022).

Functional acceleration adds momentum directly in policy space rather than in parameter space. Working in the dual variables $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$3, the paper derives extra-gradient, correction, and lazy-momentum PMD updates. Exact PMD with lookahead or extra-gradient achieves a $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$4-rate, and the construction is independent of the policy parameterization, so previous momentum schemes in parameter space appear as special cases (Chelu et al., 2024).

The choice of mirror map is itself an algorithmic degree of freedom. Using $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$5-potential mirror maps and evolutionary search, one can learn mirror geometries that outperform the conventional negative-entropy choice in Grid-World and MinAtar. These learned mirror maps transfer across several environments, while aggressive entropy-like geometries do not uniformly dominate. This suggests that PMD is better understood as a family of geometry-dependent policy-improvement operators rather than as a single KL-based method (Alfano et al., 2024).

6. Applications, empirical regularization studies, and open problems

PMD has been exported beyond single-agent tabular control. In cooperative multi-agent RL, Heterogeneous-Agent Mirror Descent Policy Optimization uses the multi-agent advantage decomposition lemma and sequential per-agent KL-regularized mirror steps under CTDE. It handles heterogeneous policy classes, including categorical, Gaussian, and Beta policies, and is reported to outperform HATRPO and HAPPO on Multi-Agent MuJoCo and SMAC $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$6 (Nasiri et al., 2023).

In LLM post-training, PMD appears as KL-regularized policy improvement over sequence distributions. PMD-mean replaces the intractable log-partition term by the mean reward under the sampling policy and performs regression in log-policy space. Its population solution is equivalent to a mirror-descent subproblem with an adaptive mixed KL–$A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$7 regularizer, which makes updates more conservative when expected rewards are low. On math reasoning tasks, this method improves stability and time efficiency, and the paper reports higher AIME performance than GRPO and strong robustness relative to PMD-part (Xu et al., 5 Feb 2026).

In continuous-control RL with expressive generative policies, One-Step Flow Policy Mirror Descent instantiates the KL-regularized PMD target with flow-based policies and shows that one-step sampling is justified by a variance–discretization bound for straight interpolation flow matching. On MuJoCo benchmarks, the method attains performance comparable to diffusion-policy baselines while requiring hundreds of times fewer function evaluations at inference (Chen et al., 31 Jul 2025).

Large-scale empirical work has also clarified PMD’s regularization landscape. A study of more than $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$8 training seeds on small environments examined the interaction between the trust-region “Drift” regularizer and the MDP regularizer. The two can partially substitute for one another, but the pairing $A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$9 and the temperatures $$d^\pi_{\rho}(s)=(1-\gamma)\sum_{t=0}^\infty \gamma^t \mathbb{P}^\pi(s_t=s\mid s_0\sim \rho).$$0 are critical; performance heat maps show an “L”-shaped region of good performance, constant $$d^\pi_{\rho}(s)=(1-\gamma)\sum_{t=0}^\infty \gamma^t \mathbb{P}^\pi(s_t=s\mid s_0\sim \rho).$$1 performs markedly better than annealing, and robustness varies sharply across mirror-map and distance choices (Kleuker et al., 11 Jul 2025).

Several open problems recur across the literature. One is statistical: for inexact PMD under a generative model, the known $$d^\pi_{\rho}(s)=(1-\gamma)\sum_{t=0}^\infty \gamma^t \mathbb{P}^\pi(s_t=s\mid s_0\sim \rho).$$2 dependence remains far from the minimax $$d^\pi_{\rho}(s)=(1-\gamma)\sum_{t=0}^\infty \gamma^t \mathbb{P}^\pi(s_t=s\mid s_0\sim \rho).$$3 lower bound (Johnson et al., 2023). Another is geometric: step-size necessity is established for natural policy gradient, but extending necessity results to general mirror maps remains open (Johnson et al., 2023). A third is representational: broad policy-convergence and finite-step results are tabular, and extensions to function approximation, continuous actions, and sampling noise remain technically delicate (Lin et al., 2022, Lan et al., 2022). A final misconception addressed by recent work is that explicit action-space exploration is always indispensable; under uniform mixing and non-diminishing minimal visitation, stochastic PMD can obtain $$d^\pi_{\rho}(s)=(1-\gamma)\sum_{t=0}^\infty \gamma^t \mathbb{P}^\pi(s_t=s\mid s_0\sim \rho).$$4 sample complexity without incorporating explicit exploration strategies (Li et al., 2023).