Technical Summary: Revisiting Policy Gradients for Restricted Policy Classes

Problem Statement

Standard policy gradient (PG) methods are widely used to solve Markov Decision Processes (MDPs). While they are theoretically guaranteed to converge to the optimal policy in unrestricted policy classes under sufficient exploration, they frequently fail in restricted policy classes. In these settings, which include state aggregation, decentralized multi-agent systems, and partially observable environments, PG methods often get stuck in suboptimal critical points.

The authors identify the root cause of this failure as the myopic nature of the standard policy gradient theorem. The theorem relies on a one-step Q-function, which evaluates the improvement of a policy based on a single action step before reverting to the current policy. In restricted classes, this "short-sightedness" means that even if a better policy exists within the class, the gradient may indicate no improvement direction because the immediate one-step cost of deviating from the current policy appears high, while the long-term benefits (which require multiple steps of the new policy) are not captured in the gradient calculation. This issue persists even in fully observable settings if the initial state distribution lacks sufficient exploration to cover the support of the optimal policy.

Methodology

1. Correlated Policy Formulation

To address myopicness, the authors propose a reformulation of the policy. Instead of viewing a stochastic policy as a direct mapping $\pi: S \to \Delta(A)$, they define a Correlated Policy $\tilde{\pi}$ as a distribution over a set of deterministic policies $\Pi_{det}$.

• Restricted Class: The learning process operates over $\tilde{\Pi}_{res} = \Delta(\Pi_{det})$, the set of all distributions over a pre-defined set of deterministic policies.
• Execution: In a standard rollout, a deterministic policy is sampled from $\tilde{\pi}$ at every step independently. However, the authors introduce a k-step evaluation mechanism.

2. k-step Policy Gradients

The core innovation is the k-step Policy Gradient, which modifies the evaluation horizon to escape local optima.

• k-step Value Function ($J_{\tilde{\pi}, k}$): Instead of sampling a new deterministic policy at every step, the algorithm samples a deterministic policy $\pi_{det} \sim \tilde{\pi}$ and executes it for $k$ consecutive steps before resampling.
• k-step Q-function ($Q_{\tilde{\pi}, k}$): This function evaluates the return of taking a specific deterministic policy $\pi'_{det}$ for the first $k$ steps, followed by the standard k-step rollout of $\tilde{\pi}$.
• Theorem 5.5 (k-step PG Theorem): The gradient of the k-step value function with respect to the parameters of $\tilde{\pi}$ is expressed using the k-step Q-function:
  $$\nabla_\alpha J_{\tilde{\pi}_\alpha, k}(\mu) = \frac{1}{1-\gamma^k} \mathbb{E}_{s \sim d_{\tilde{\pi}_\alpha, k}^\mu} \left[ \mathbb{E}_{\pi_{det} \sim \tilde{\pi}_\alpha} [Q_{\tilde{\pi}_\alpha, k}(s, \pi_{det}) \nabla_\alpha \log \tilde{\pi}_\alpha(\pi_{det})] \right]$$
  By evaluating actions over a $k$-step window, the gradient can "see" the long-term benefits of a deterministic policy that might incur short-term costs, thereby correcting the myopic bias.

3. Optimization Algorithms

The paper analyzes the convergence of Projected Gradient Descent (PGD) and Mirror Descent when applied to this k-step objective. The analysis assumes only the smoothness and differentiability of the value function, avoiding strong assumptions like Bellman completeness or gradient dominance that are common in prior literature.

Key Contributions and Results

1. Theoretical Guarantee of Near-Optimality

The primary theoretical contribution is the proof that all zero-gradient critical points and boundary local minima of the k-step objective are exponentially close to the optimal deterministic policy in terms of performance.

• Theorem 6.1: For any critical point $\tilde{\pi}_{crit}$, the performance gap satisfies:
  $$\mathbb{E}_{\pi_{det} \sim \tilde{\pi}_{crit}}[J_{\pi_{det}}(\mu)] - J_{\pi^*_{det}}(\mu) \leq \frac{8\gamma^k}{1-\gamma} g_{max}$$
  As $k$ increases, the bound tightens exponentially, effectively eliminating suboptimal local minima that plague 1-step methods.

2. Convergence Rates

The authors demonstrate that descent algorithms (PGD and Mirror Descent) converge to these near-optimal solutions at a rate of $O(1/T)$.

• Theorem 6.2 & 6.3: Despite the absence of the standard gradient dominance condition (which often fails in restricted classes), the k-step evaluation induces an approximate gradient dominance condition. This allows the algorithms to achieve the faster $O(1/T)$ rate, typically associated with convex or gradient-dominant problems, rather than the slower $O(1/\sqrt{T})$ rate typical for smooth non-convex functions.

3. Avoidance of Distribution Mismatch

A significant finding is that these bounds do not depend on the distribution mismatch factors $\|d_{\pi^*}^\mu / d_{\pi}^\mu\|_\infty$ or $\|d_{\pi^*}^\mu / \mu\|_\infty$.

• In standard PG theory, these factors can be infinite if the initial distribution $\mu$ does not cover the support of the optimal policy's state occupancy, leading to suboptimal critical points even in fully observable settings.
• The k-step method is shown to escape these points regardless of the initial distribution's coverage, provided $k$ is sufficiently large.

Significance and Applications

The paper claims that this work provides a theoretical foundation for solving complex MDPs with restricted policy classes where previous methods failed or relied on unrealistic assumptions.

• State Aggregation: The method offers near-optimal guarantees for policies that map aggregated states to actions, a common necessity in large-scale systems.
• Multi-Agent Systems: It applies to cooperative settings with independent agents, decentralized agents, and group-decentralized agents. These are often modeled as Dec-POMDPs, which are known to be intractable (NEXP-complete). The paper provides a positive theoretical result for gradient-based methods in these settings.
• Fully Observable Settings: The work highlights that even in fully observable MDPs, poor exploration (initial state distribution) can create suboptimal critical points. The k-step method provides a mechanism to escape these without requiring perfect initial exploration.

Conclusion

By shifting from a myopic 1-step evaluation to a correlated k-step evaluation, the proposed method fundamentally alters the optimization landscape of restricted policy classes. It guarantees that critical points are exponentially close to the optimal solution and ensures fast convergence for standard descent algorithms, all without requiring strong realizability assumptions or perfect initial exploration. This bridges a gap between theoretical guarantees and practical applications in state aggregation and multi-agent reinforcement learning.