🤖 machine learning

Global linear convergence of entropy-regularized softmax policy gradient beyond tabular MDPs

This paper establishes the global linear convergence of entropy-regularized softmax policy gradient with log-linear function approximation for infinite-horizon MDPs with continuous state and action spaces by proving a non-uniform Polyak-Łojasiewicz inequality under specific feature regimes that ensure the Fisher information matrix or uncentered covariance matrix remains well-conditioned.

Original authors: Ziyue Chen, David Šiška, Lukasz Szpruch

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ziyue Chen, David Šiška, Lukasz Szpruch

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to play a complex video game. The robot has to make decisions (actions) based on what it sees (states) to get the highest score. In the world of Reinforcement Learning (RL), this is called finding the "optimal policy."

For a long time, mathematicians could only prove that the robot would learn quickly and reliably if the game was very simple—like a board game with a fixed number of squares and moves. This is called the "tabular" setting. But real life is messy; the state space is continuous (like driving a car where speed and position can be any number), and the actions are infinite.

This paper by Chen, Šiška, and Szpruch tackles the hard question: Can we prove that a robot learns efficiently in these complex, continuous worlds if we use a specific type of "smart" learning algorithm?

Here is the breakdown of their findings using everyday analogies.

1. The Problem: The "Hilly" Landscape

Imagine the robot's goal is to find the highest peak in a vast, foggy mountain range. The "height" of the mountain represents how good the robot's strategy is.

The Challenge: In many learning algorithms, the mountain range is full of false peaks (local optima). The robot might get stuck on a small hill thinking it's the top, never reaching the real summit.
The Twist: The authors add a special ingredient called Entropy Regularization. Think of this as a "curiosity bonus." The robot is rewarded not just for getting a high score, but for keeping its options open and not being too rigid. Mathematically, this smooths out the mountain range, making it easier to find the true peak.

2. The Method: The "Log-Linear" Map

Since the mountain is too big to map every single inch (the continuous state space), the robot uses a simplified map.

The Analogy: Instead of memorizing every tree and rock, the robot uses a set of "features" (like "is it steep?", "is it sunny?", "is there a river?"). It combines these features using a linear formula (a weighted sum) to decide what to do. This is called Log-Linear Softmax Policy.
The Goal: The authors want to prove that if the robot follows the "gradient flow" (a mathematical way of saying "always walk uphill"), it will reach the top of the mountain exponentially fast. This means it doesn't just get better slowly; it gets better at a speed that doubles its progress every second.

3. The Big Hurdle: The "Slippery Slope"

In the simple "tabular" world, the math is nice and round. But in this complex world, the shape of the mountain changes depending on where you are.

The Issue: Sometimes, the ground gets so flat or slippery that the robot might stop moving or move incredibly slowly. In math terms, the "Fisher Information Matrix" (a measure of how much information the robot's current view gives it) can become "degenerate" or lose its grip.
The Paper's Solution: The authors prove a Non-Uniform Polyak–Łojasiewicz (PŁ) Inequality.
- Simple Translation: They proved that even though the ground is slippery in some spots, the "pull" toward the top is always strong enough to keep the robot moving, provided the robot doesn't get stuck in a specific weird configuration.

4. The Secret Sauce: Two Types of "Maps"

To guarantee the robot never gets stuck, the authors identified two specific types of "feature maps" (the way the robot sees the world) that work perfectly.

Type A: The "Full Affine Span" (The Trigonometric Map)

The Analogy: Imagine the robot uses a map based on waves (sine and cosine waves), like the Fourier basis.
Why it works: The authors proved that with this map, if the robot tries to go too far in any direction, the "curiosity bonus" (Entropy) becomes infinitely large. It's like a rubber band that gets infinitely tight if you stretch it too far. This forces the robot to stay within a safe, bounded area where the ground is never too slippery.
Result: The robot is guaranteed to find the peak quickly.

Type B: The "Simplex" Features (The Bernstein Map)

The Analogy: Imagine the robot uses a map based on probability percentages (like the Bernstein polynomials), where all the weights must add up to 100%.
The Nuance: In this case, the "rubber band" (Entropy) only gets tight if the robot tries to stretch in a specific direction (perpendicular to the "all-equal" direction).
Result: Even with this slightly different map, the authors proved the robot still stays in a safe zone and converges to the peak linearly.

5. What They Proved (The Bottom Line)

The paper provides a rigorous mathematical guarantee:

Global Convergence: The robot will eventually find the best possible strategy, no matter where it starts.
Linear Speed: It won't just get there; it will get there fast, with the error shrinking by a constant percentage every step (like compound interest, but in reverse).
Beyond Simple Games: This works for complex, continuous environments, not just simple grids.

What They Did NOT Claim

It is important to stick to what the paper actually says:

They did not claim this works for every possible type of feature map. They specifically identified the "Full Affine Span" and "Simplex" types.
They did not claim this solves the problem of "approximation error" (where the map itself is a bad approximation of reality). They assumed the "Q-realizability" condition, meaning the true optimal strategy can be represented by their chosen map.
They did not discuss clinical uses, self-driving cars, or specific video games. They focused purely on the theoretical convergence of the algorithm in a mathematical model.

In summary: The authors took a difficult, continuous learning problem and showed that if you use the right kind of "features" (maps) and add a "curiosity bonus," the learning algorithm is mathematically guaranteed to zoom straight to the best solution without getting stuck.

Technical Summary: Global Linear Convergence of Entropy-Regularized Softmax Policy Gradient Beyond Tabular MDPs

Problem Statement
This paper investigates the global convergence properties of policy gradient (PG) methods for infinite-horizon entropy-regularized Markov Decision Processes (MDPs) with continuous state and action spaces. While entropy regularization (introducing a KL-divergence term with intensity $\tau > 0$ ) is known to improve the convergence landscape of softmax policies in tabular settings (finite $S$ and $A$ ), extending these guarantees to continuous spaces with function approximation remains a significant theoretical challenge.

Specifically, the authors consider log-linear softmax policies defined by:
$\pi_\theta(da|s) \propto \exp(\langle \theta, g(s, a) \rangle) \mu(da)$
where $g: S \times A \to \mathbb{R}^p$ represents a feature basis, $\theta \in \mathbb{R}^p$ are the parameters, and $\mu$ is a reference measure. The objective is to minimize the regularized value function $V^\pi_\tau(\rho)$ . The core difficulty lies in the non-convexity of the objective with respect to $\theta$ , even with entropy regularization, and the potential degeneracy of the Fisher Information Matrix (FIM) or feature covariance matrices in continuous settings, which can prevent standard convergence proofs from holding.

Methodology
The analysis relies on a continuous-time gradient flow formulation:
$\frac{d}{dt}\theta_t = -\nabla_\theta V^{\pi_{\theta_t}}_\tau(\rho)$
The proof strategy proceeds in three main stages:

Non-Uniform Polyak–Łojasiewicz (PŁ) Inequality:
The authors first establish a non-uniform PŁ inequality of the form:
$V^{\pi_\theta}_\tau(\rho) - V^{\pi^*_\tau}_\tau(\rho) \leq C(\theta) \|\nabla_\theta V^{\pi_\theta}_\tau(\rho)\|^2$
This inequality links the suboptimality gap to the squared gradient norm. Crucially, the constant $C(\theta)$ depends on the policy geometry, specifically the smallest eigenvalue of the FIM (or an uncentered feature covariance matrix for simplex features) and the density bounds of the policy. Unlike the tabular case where $C$ can be uniform, here $C(\theta)$ is non-uniform and could theoretically degenerate as $\|\theta\| \to \infty$ .
Q $\pi_\tau$ -Realizability Assumption:
The analysis assumes that for any policy $\pi$ , the regularized state-action value function $Q^\pi_\tau$ can be exactly represented by the linear feature basis:
$\langle \theta(\pi), g(s, a) \rangle = -\frac{1}{\tau} Q^\pi_\tau(s, a)$
This assumption ensures that the optimal policy $\pi^*_\tau$ lies within the parameterized class, i.e., $\pi^*_\tau = \pi_{\theta^*}$ for some $\theta^*$ .
Lyapunov Analysis and Radial Unboundedness:
To prove global linear convergence, the authors must show that the non-uniform constant $C(\theta_t)$ remains bounded along the gradient flow trajectory. This requires proving that the trajectory $\{\theta_t\}$ remains in a compact set.
- They utilize the regularized objective $V^{\pi_\theta}_\tau(\rho)$ as a Lyapunov function, noting that it decreases monotonically.
- They prove that under specific conditions on the feature basis $g$ , the KL-divergence term $KL(\pi_\theta | \mu)$ is radially unbounded (i.e., $KL \to \infty$ as $\|\theta\| \to \infty$ ). This property prevents the parameters from diverging to infinity while the value function remains bounded, thereby ensuring the trajectory stays within a compact region where the eigenvalues of the relevant matrices are bounded away from zero.

Key Contributions and Results

The paper identifies two distinct regimes for feature bases under which global linear convergence is achieved:

Full-Affine-Span Features:
- Assumption: The feature map $g(s, \cdot)$ has full affine dimension, and the set of maximizers of any linear projection of $g$ has zero measure under $\mu$ .
- Result: The KL divergence is radially unbounded in all directions. Consequently, the smallest eigenvalue of the FIM remains bounded below by a positive constant dependent on initialization.
- Convergence: The suboptimality decays as $O(e^{-Ct})$ for some $C > 0$ .
Simplex-Valued Features:
- Assumption: The features lie in the probability simplex ( $\sum g_i = 1, g_i \geq 0$ ). The condition on the maximizer set is relaxed to apply only to vectors orthogonal to the all-ones vector.
- Result: The KL divergence is radially unbounded only in the subspace orthogonal to the all-ones vector (the direction of parameter redundancy). The authors prove that the gradient flow preserves the projection of $\theta$ onto the all-ones vector, effectively confining the dynamics to the relevant subspace.
- Convergence: Using an uncentered feature covariance matrix instead of the FIM, they establish a uniform lower bound on the eigenvalues, yielding the same global linear convergence rate $O(e^{-Ct})$ .

Significance and Claims
The authors claim that this work extends the global convergence theory of entropy-regularized softmax policy gradient beyond the tabular setting established by Agarwal et al. (2020), Bhandari and Russo (2024), and Mei et al. (2020).

Novelty in Continuous Settings: The paper addresses the computational intractability of tabular methods in continuous spaces by rigorously analyzing linear function approximation.
Handling Non-Convexity: By establishing a non-uniform PŁ inequality and controlling the non-uniform constant via Lyapunov techniques, the authors overcome the non-convexity of the softmax parameterization in continuous spaces.
Role of KL Regularization: The paper highlights a novel use of the KL divergence as a Lyapunov function to ensure radial unboundedness, a property that does not hold in the standard tabular case (where the KL term is finite for any policy with full support).
Feature Sensitivity: The work provides specific examples (e.g., trigonometric features, Bernstein polynomials) that satisfy the convergence conditions and counter-examples (e.g., hat functions) where radial unboundedness fails, clarifying the necessary geometric properties of feature bases for linear convergence.

The results imply that, under $Q^\pi_\tau$ -realizability and suitable feature conditions, entropy-regularized policy gradient converges globally and linearly to the optimal policy in general state and action spaces, matching the favorable rates previously known only for tabular or linear-quadratic settings.