Rethinking the Global Convergence of Softmax Policy Gradient with Linear Function Approximation: The Case of Multi-Armed Bandits

Max Qiushi Lin; Jincheng Mei; Matin Aghaei; Michael Lu; Bo Dai; Alekh Agarwal; Dale Schuurmans; Csaba Szepesvari

cs.LGMay 6, 2025

Rethinking the Global Convergence of Softmax Policy Gradient with Linear Function Approximation: The Case of Multi-Armed Bandits

Authors: Max Qiushi Lin, Jincheng Mei, Matin Aghaei, Michael Lu, Bo Dai, Alekh Agarwal, Dale Schuurmans, Csaba Szepesvari, +1 more

Organizations: Simon Fraser University · Google DeepMind · Google Research · Google DeepMind and University of Alberta

Abstract

Policy gradient (PG) methods have played an essential role in the empirical successes of reinforcement learning. In order to handle large state-action spaces, PG methods are typically used with function approximation. In this setting, the approximation error in modeling problem-dependent quantities is a key notion for characterizing the global convergence of PG methods. We study Softmax PG with linear function approximation (referred to as $\texttt{Lin-SPG}$ ) and demonstrate that the approximation error is irrelevant to the algorithm's global convergence even in the bandit setting. Consequently, we rethink the effect of approximation error in the standard stochastic multi-armed bandit problem. We first identify the conditions on the policy feature representation that can guarantee the asymptotic global convergence of $\texttt{Lin-SPG}$ . Under these feature conditions, we further prove that $T$ iterations of $\texttt{Lin-SPG}$ with a problem-specific learning rate result in an $O(1/T)$ convergence to the optimal policy. Moreover, we prove that $\texttt{Lin-SPG}$ with an arbitrary constant learning rate can ensure asymptotic convergence to the optimal policy.

Figures & tables

Figure 1: Visualization of the optimization landscape for Examples 1 , 2 , 3 and 4 . For each $\theta\in{\mathbb{R}}^{2}$ , we calculate the expected reward ${\left\langle\pi_{\theta},r\right\rangle}$ for the log-linear policy $\pi_{\theta}$ defined in Eq. 3 and color the optimization landscape with respect to its value. The white arrows indicate the gradient directions, whereas the red arrows demonstrate the optimization trajectories of running Algorithm 1 in these four examples. Specifically, in both Examples 1 and 2 , we set $\theta_{1}=(3,3)^{\top}$ , while in both Examples 3 and 4 , we set $\theta_{1}=(-3,-1.2)^{\top}$ . The learning rate $\eta$ is set to $0.2$ for all examples. Intuitively, for Examples 1 and 3 , all the gradient directions in the landscape demonstrate that Algorithm 1 can converge to the optimal policy starting from any initialization. On the other hand, for Examples 2 and 4 , we can identify a specific region of the landscape from which Algorithm 1 fails to converge globally.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 2 — Figure 2: The effect of feature conditions on the global convergence.

Panel of Figure 3 — Figure 3: Convergence of Lin-SPG in the exact setting for (a) $K=3$ and (b) $K=6$ across different feature dimensions.

Figure 4: Convergence of Lin-SPG in the stochastic setting for $K=6$ and $d=3$ under different constant learning rates across different reward distributions.

Explore similar work

Global linear convergence of entropy-regularized softmax policy gradient beyond tabular MDPs
May 24, 2026Ziyue Chen, David Šiška, Lukasz SzpruchEntropy Regularized Reinforcement Learning Policy Gradient
Vanishing L2 regularization for the softmax Multi Armed Bandit
May 5, 2026Stefana-Lucia Anita, Gabriel TuriniciEntropy Regularized Reinforcement Learning Policy Gradient
CEREMADE, Universit´e Paris Dauphine - PSL, CNRS, Paris, France · ‘Octav Mayer” Institute of Mathematics of the Romanian Academy, Bd. Carol I 8, Ias¸i 700505, Romania
Delightful Gradients Accelerate Corner Escape
May 12, 2026Jincheng Mei, Ian OsbandPolicy Gradient Optimal Policies
Google DeepMind

May 24, 2026, cs.LGJ/K move · Enter open · S save

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

Ziyue Chen, David Šiška, Lukasz Szpruch

Entropy Regularized Reinforcement Learning Policy Gradient Convergence Markov Decision Processes Value Functions

Open paper AskPDF arXiv

We study the global convergence of policy gradient for infinite-horizon entropy-regularized Markov decision processes (MDPs) with continuous state and action spaces. We consider log-linear softmax policies with linear function approximation, which extend the tabular softmax parameterization while retaining a tractable policy class. Under

Q^π_τ

-realizability for the regularized state-action value function, we first establish a non-uniform Polyak--Łojasiewicz (PŁ) inequality. The non-uniformity arises through degeneracy of constants associated with the policy geometry, namely the Fisher information matrix or an uncentered feature covariance matrix. We then identify two feature regimes under which this non-uniform constant can be bounded along the gradient flow. For full-affine-span features, we prove radial unboundedness of the KL regularizer and show that the smallest eigenvalue of the Fisher information matrix remains bounded below by an initialization-dependent positive constant. For simplex-valued features, we prove an analogous radial unboundedness result in the subspace orthogonal to the all-ones vector and obtain a uniform lower bound for the smallest eigenvalue of the uncentered covariance matrix. These results imply global linear convergence of the regularized objective along the gradient flow, i.e. suboptimality decaying as

\mathcal{O}(e^{-Ct})

for some

C>0

. Our analysis extends the global convergence theory of entropy-regularized softmax policy gradient beyond the tabular setting of Agarwal et al. (2020); Bhandari and Russo (2024); Mei et al. (2020).

Abstract

Explore similar work

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

Vanishing L2 regularization for the softmax Multi Armed Bandit

Delightful Gradients Accelerate Corner Escape