On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift 4: Log-Barrier-Regularized Softmax Policy Gradient Is ε-Optimal After 320|S|²|A|²D∞²/((1−γ)⁶ε²) StepsResearch Paper
Motivation
Policy gradient methods search directly over parameterized policies of a Markov decision process (MDP) by gradient ascent on the expected discounted reward. They are among the most widely used methods in reinforcement learning, yet the value is a nonconcave function of the parameters, and until recently little was known about whether, or how fast, such methods reach a globally optimal policy. Agarwal, Kakade, Lee and Mahajan (arXiv:1908.00261v5, JMLR 22(98), 2021) give a systematic answer for tabular and function-approximation settings.
This mission formalizes one of their central results. Under the softmax parameterization, plain gradient ascent converges to optimal values only asymptotically (their Theorem 5.1, rate unknown): parameters can diverge and policies become nearly deterministic early, which stalls progress. Adding a log barrier regularizer, equivalent to a relative-entropy penalty towards the uniform policy, keeps action probabilities away from zero, and the paper shows that this yields an iteration count polynomial in every problem quantity. The result is the "softmax + log barrier" row of the paper's Table 1 and a reference point for later analyses of entropy-regularized policy gradient (for example Mei et al., 2020; Cen et al., 2021).
Setting
A finite MDP consists of finite sets (states) and (actions, nonempty), a transition kernel , rewards and a discount factor . A policy assigns a probability distribution over actions to every state. Its value is , and for a distribution over states . The action value is and the advantage is . An optimal policy maximizes at every state simultaneously; .
The discounted state visitation distribution is . Given a start distribution used by the algorithm, the distribution mismatch coefficient is .
The softmax policy with parameter is . For the log barrier regularized objective is
and gradient ascent with step size is . Write .
Formalization targets
Goal: Corollary 5.1 (iteration complexity)
Let for all (for example the mismatch coefficient), , , and . Starting from a parameter whose softmax policy is uniform (e.g. ),
Milestones
- Lemma 3.2 (performance difference): .
- Lemma C.1 (softmax gradient): .
- Display (14): the same plus for .
- Opening claim of the proof of Theorem 5.2: at an -stationary point with , .
- Theorem 5.2: under the same hypotheses, .
- Lemma D.4: has -Lipschitz gradient, .
- Theorem E.1(3) (Beck, Theorem 10.15), unconstrained ascent form: for -smooth and step , .
Significance
The corollary shows that a first-order method on a nonconcave objective reaches an -optimal policy in a number of steps polynomial in , , , and the mismatch coefficient, with no dependence on the MDP's structure beyond these quantities. Theorem 5.2 is of independent use: it certifies any approximately stationary point of the regularized objective, however it was found, as approximately optimal for the original problem. The dependence on makes precise why the exploratory start distribution matters.
The results are proved in the paper; none of them is machine-checked. A formal development adds a checked softmax policy gradient formula on an infinite-horizon discounted MDP, a checked smoothness bound for the regularized objective, the nonconvex gradient-method rate in Euclidean space, and a corrected statement of the corollary (see the next sections). Lemma 3.2 and Lemma C.1 are reused across the other missions of this series.
Difficulty
The value function is not concave in and its gradient vanishes where policies become deterministic, so the standard route "small gradient implies near-optimal" fails for the unregularized objective. With the regularizer, the link between stationarity and optimality has to hold with explicit constants that do not degrade as policies approach the boundary of the simplex, which is where the unregularized argument breaks down. The second obstacle is smoothness: is an infinite discounted sum of products of transition probabilities, and its Hessian must be bounded uniformly in by , which in Lean means differentiating a tsum of matrix-power–like terms in . The policy gradient formula itself requires the same interchange of differentiation and infinite summation.
Formalization scope
States and actions are finite types; is nonempty. Parameters live in EuclideanSpace ℝ (S × A), so is the norm and is Mathlib's gradient. The MDP layer (transition kernel, policy, occupation probabilities, , , optimal policy) is the published FoundationsML.ReinforcementLearning library, and , , are defined on top of it. Standing assumptions are those of §3: a kernel, , , probability distributions, optimal at every state.
Conventions committed to:
- the mismatch coefficient is replaced by any with componentwise; no division by occurs, and the same sets ;
- is "there is with …";
- , implicit in the paper, is a hypothesis of Theorem 5.2 and its proof-step milestone; in Lemma D.4;
- Theorem E.1 is posed for ascent on with an upper bound in place of an attained optimum.
Two corrections of the printed corollary. The printed statement allows "any initial ", but its proof uses , valid for a uniform initial policy and false in general; the formal goal assumes a uniform start. The printed contradicts the proof, which takes from Lemma D.4, ; the formal goal uses the latter, with the same iteration count.
The goal is not trivialized by its encoding: the hypotheses fix the run completely from , and , they are jointly satisfiable (a sorry-free check exhibits a one-state, two-action MDP with an optimal policy and ), and the conclusion concerns the unregularized value of the iterates. Theorem 5.2, Lemma D.4 and Theorem E.1 are milestones, not hypotheses of the goal.
Infrastructure that a complete development needs, and that is reusable beyond this mission: differentiability of discounted values in the policy parameters and the policy gradient theorem for the softmax class; Hessian bounds for softmax; the descent lemma and the rate of gradient methods for smooth nonconvex functions on Euclidean space. Contributions to any of these, or proofs of single milestones, are welcome.
Selected references
- A. Agarwal, S. M. Kakade, J. D. Lee, G. Mahajan, On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift, JMLR 22(98), 2021; arXiv:1908.00261v5. https://arxiv.org/abs/1908.00261v5
- S. Kakade, J. Langford, Approximately Optimal Approximate Reinforcement Learning, ICML 2002. https://dl.acm.org/doi/10.5555/645531.656005
- A. Beck, First-Order Methods in Optimization, SIAM, 2017. https://doi.org/10.1137/1.9781611974997
- J. Mei, C. Xiao, C. Szepesvári, D. Schuurmans, On the Global Convergence Rates of Softmax Policy Gradient Methods, ICML 2020. https://arxiv.org/abs/2005.06392
- S. Cen, C. Cheng, Y. Chen, Y. Wei, Y. Chi, Fast Global Convergence of Natural Policy Gradient Methods with Entropy Regularization, Operations Research 70(4), 2022. https://arxiv.org/abs/2007.06558