Global Optimality Guarantees for Policy Gradient Methods 4: In Finite-Horizon Problems with a Non-Stationary Policy Class Containing an Optimal Policy, Every Stationary Point Is OptimalResearch Paper
Motivation
Policy gradient methods optimize a parameterized policy by gradient descent on its expected total cost. They are a staple of reinforcement learning and of simulation-based optimization in operations research, but the objective is nonconvex in the policy parameters, and the standard theory guarantees convergence only to a stationary point. J. Bhandari and D. Russo (arXiv:1906.01786v3, Operations Research, 2024, DOI 10.1287/opre.2021.0014) identify structural conditions under which every stationary point is globally optimal.
Their first result (Theorem 1) requires the policy class to be closed under policy improvement. That excludes many classes used in practice, where it is known only that some member of the class is optimal: base-stock policies in finite-horizon inventory control are the standard example (Bertsekas, Dynamic Programming and Optimal Control, 1995). For stationary policy classes, containing an optimal policy is not enough: the paper's Example 1 has a policy class that contains the optimal policy, and policy gradient still gets stuck in a bad local minimum. Theorem 3 shows that for finite-horizon problems with a non-stationary policy class (separate parameters for each period), containing an optimal policy suffices, together with a single-period regularity condition imposed only at the optimal cost-to-go. This mission formalizes Theorem 3.
Setting
A Markov decision process consists of a state space , nonempty feasible action sets , a bounded measurable cost , a transition kernel , a discount factor and an initial distribution . A policy is a measurable map with ; is the set of all of them. The cost-to-go is , the loss is , and the discounted state-occupancy measure is . An optimal policy has pointwise for all ; . With the Bellman operator , the weighted policy-iteration objective is
A parameterized policy class is for a convex , and . A point is stationary for if is differentiable at and for all .
The finite-horizon structure (Condition 3) embeds a time-inhomogeneous problem in this discounted formulation: the states split into stages ; from a state in , , every feasible action leads to ; and is a costless absorbing state. The parameter is a concatenation , the parameter set is a product , and on the action depends only on . The stage objectives are .
Condition 4 asks that for every , the problem has no suboptimal stationary points. Assumption 3 asks that for every . Condition 0 asks that the policy-iteration objective be continuously differentiable in the parameters.
Formalization targets
Goal: Theorem 3 (p. 17)
Under Conditions 0, 3 and 4 and Assumption 3, if is an optimal policy for some , then every stationary point of on satisfies
Milestones
- Lemma 6 (policy gradient theorem): .
- Lemma 17 (balance equation): .
- The stage-wise balance equation: for , .
- Separability (30): minimizes over iff each minimizes over .
- Stage-wise stationarity (31): is stationary for iff each is stationary for on .
- Lemma 1: iff -almost surely.
Significance
Theorem 3 covers finite-horizon dynamic programs in which a structured class of time-varying policies is known to contain an optimal policy, but is not closed under policy improvement: the paper's Example 8, multi-period inventory control with base-stock policies, is the motivating case. Policy gradient on such a class has no suboptimal stationary points, so any method that finds a stationary point finds a globally optimal policy, optimal among all policies and not only within the class. Condition 4 involves only the Bellman objective of the optimal cost-to-go , which often has regularity (for instance convexity) that the cost-to-go of an arbitrary policy lacks.
The result is proved in the paper. To our knowledge no part of it, nor the policy gradient theorem for deterministic policies on general state spaces, has a machine-checked proof. Formalizing it requires a measure-theoretic treatment of occupancy measures, the stage-wise decomposition of the weighted policy-iteration objective, and the backward induction of App. D.2.
Difficulty
The obvious approach reasons stage by stage directly on , but couples the stages: changing changes the distribution of states visited after stage . The paper's argument instead passes through the policy gradient theorem, which turns stationarity of into stationarity of single-period problems weighted by the occupancy measure of the current policy. These single-period problems involve , while Condition 4 speaks about ; relating them requires a backward induction over stages in which almost-sure equality on stage is transported to stage through the balance equation. This step needs the occupancy measures of all policies in the class, including policies other than , to be equivalent to , which is where Assumption 3 enters. The naive variant with a stationary policy class fails (Example 1).
Formalization scope
The state and action spaces are arbitrary measurable spaces (the paper takes Borel subsets of Euclidean spaces; no argument uses that structure). Cost and kernel are given on all state-action pairs; the cost is bounded. Policies are deterministic, stationary and measurable; parameters live in EuclideanSpace ℝ (Fin d), so gradients and norms are Euclidean. is never defined as an infimum: an optimal policy is a hypothesis, with . Stages and sub-vectors are numbered from by a measurable stage index and a block index on coordinates; the product structure of is encoded as closure under exchanging one block.
Hypotheses beyond the printed sentence of Theorem 3, all disclosed in the item statements: Assumption 3 (announced on p. 17 as required, used in the proof); Condition 0 (needed for Lemma 6, which the proof invokes), in a joint form that the proof of Lemma 6 uses and that implies the printed partial-map form; the standing Assumptions 1 and 2 of §2; and, inside Condition 4, differentiability of on , the premise of Definition 1. The Bellman operators carry the factor omitted in the printed (3)–(4). The stage-wise balance equation is stated for ; the printed range fails at the absorbing state.
A formalization in which the goal assumes the stage-wise characterization (31), the balance equation, or on any stage is not Theorem 3; neither is one where "contains an optimal policy" means optimal only within . Stationarity includes differentiability, so a junk zero gradient never makes a point stationary.
Needed infrastructure: occupancy measures as series of iterated kernels and their balance equation, the cost-to-go as a bounded measurable function, the performance difference identity, the policy gradient theorem, and the decomposition of integrals over the stage partition. The occupancy-measure and policy gradient results are reusable for every mission of this series; contributions of these general lemmas are welcome.
Selected references
- J. Bhandari, D. Russo, Global Optimality Guarantees for Policy Gradient Methods, arXiv:1906.01786v3, 2022; Operations Research, 2024. https://arxiv.org/abs/1906.01786 , https://doi.org/10.1287/opre.2021.0014
- S. Kakade, J. Langford, Approximately Optimal Approximate Reinforcement Learning, ICML, 2002. https://dl.acm.org/doi/10.5555/645531.656005
- I. Osband, B. Van Roy, D. Russo, Z. Wen, Deep Exploration via Randomized Value Functions, JMLR 20, 2019. https://jmlr.org/papers/v20/18-339.html
- D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, M. Riedmiller, Deterministic Policy Gradient Algorithms, ICML, 2014. https://proceedings.mlr.press/v32/silver14.html
- D. P. Bertsekas, Dynamic Programming and Optimal Control, Athena Scientific, 1995. https://www.athenasc.com/dpbook.html