Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Beyond Softmax and Entropy: Improving Convergence Guarantees of Policy Gradients by f-SoftArgmax Parameterization with Coupled Regularization

Safwan Labbi, Daniil Tiapkin, Paul Mangold, Eric Moulines

TL;DR

The paper addresses the instability and slow convergence of policy gradient methods caused by softmax parameterization. It proposes a general f-SoftArgmax parameterization coupled with the corresponding f-divergence regularizer, aligning the optimization geometry with the regularized objective. The authors prove explicit non-asymptotic last-iterate convergence guarantees for stochastic policy gradient in finite MDPs and show polynomial sample complexity for Tsallis couplings versus exponential bounds for softmax–entropy. Empirically, Tsallis-based parameterizations improve exploration and learning in noisy and sparse-reward environments, with task-dependent optimal choices for the Tsallis parameter, suggesting broad practical value and tunable performance across RL problems.

Abstract

Policy gradient methods are known to be highly sensitive to the choice of policy parameterization. In particular, the widely used softmax parameterization can induce ill-conditioned optimization landscapes and lead to exponentially slow convergence. Although this can be mitigated by preconditioning, this solution is often computationally expensive. Instead, we propose replacing the softmax with an alternative family of policy parameterizations based on the generalized f-softargmax. We further advocate coupling this parameterization with a regularizer induced by the same f-divergence, which improves the optimization landscape and ensures that the resulting regularized objective satisfies a Polyak-Lojasiewicz inequality. Leveraging this structure, we establish the first explicit non-asymptotic last-iterate convergence guarantees for stochastic policy gradient methods for finite MDPs without any form of preconditioning. We also derive sample-complexity bounds for the unregularized problem and show that f-PG, with Tsallis divergences achieves polynomial sample complexity in contrast to the exponential complexity incurred by the standard softmax parameterization.

Beyond Softmax and Entropy: Improving Convergence Guarantees of Policy Gradients by f-SoftArgmax Parameterization with Coupled Regularization

TL;DR

The paper addresses the instability and slow convergence of policy gradient methods caused by softmax parameterization. It proposes a general f-SoftArgmax parameterization coupled with the corresponding f-divergence regularizer, aligning the optimization geometry with the regularized objective. The authors prove explicit non-asymptotic last-iterate convergence guarantees for stochastic policy gradient in finite MDPs and show polynomial sample complexity for Tsallis couplings versus exponential bounds for softmax–entropy. Empirically, Tsallis-based parameterizations improve exploration and learning in noisy and sparse-reward environments, with task-dependent optimal choices for the Tsallis parameter, suggesting broad practical value and tunable performance across RL problems.

Abstract

Policy gradient methods are known to be highly sensitive to the choice of policy parameterization. In particular, the widely used softmax parameterization can induce ill-conditioned optimization landscapes and lead to exponentially slow convergence. Although this can be mitigated by preconditioning, this solution is often computationally expensive. Instead, we propose replacing the softmax with an alternative family of policy parameterizations based on the generalized f-softargmax. We further advocate coupling this parameterization with a regularizer induced by the same f-divergence, which improves the optimization landscape and ensures that the resulting regularized objective satisfies a Polyak-Lojasiewicz inequality. Leveraging this structure, we establish the first explicit non-asymptotic last-iterate convergence guarantees for stochastic policy gradient methods for finite MDPs without any form of preconditioning. We also derive sample-complexity bounds for the unregularized problem and show that f-PG, with Tsallis divergences achieves polynomial sample complexity in contrast to the exponential complexity incurred by the standard softmax parameterization.
Paper Structure (76 sections, 52 theorems, 410 equations, 4 figures, 2 tables, 1 algorithm)

This paper contains 76 sections, 52 theorems, 410 equations, 4 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Coupling Parameterization & Regularization
  4. From Non-Uniform Łojasiewicz to Polyak-Łojasiewicz.
  5. Convergence Analysis of f-PG
  6. Convergence for unregularized objective.
  7. Sample complexity for specific choices of $\boldsymbol{f}$.
  8. Experiments
  9. Conclusion
  10. Notations
  11. $f$-Divergence.
  12. Distribution of the state-action sequence.
  13. Norms.
  14. Functional and matrix forms.
  15. Smoothness of the objective
  16. ...and 61 more sections

Key Result

Lemma 3.2

Assume that, for some $\underline{\refpolicy}>0$, $f$ and $\refpolicy$ satisfy assumF:functionF and assumP:refpolicy. For any $\in \mathcal{S}$, we have where for any vector $u \in \mathbb{R}^{|\mathcal{A}|}$, $H(u) := \mathop{\mathrm{diag}}\nolimits(u) - u u^{\top}$, $\regqfunc := \rewardMDP+\gamma \kerMDP\regvaluefunc[\theta]$, and $\visitation[\rho][\theta]() := \visitation[\rho][\cpolicy{\the

Figures (4)

  • Figure 1: Regularized value landscapes (with temperature $\lambda=1$) for a one-state, two-action MDP: softmax with entropy (left) versus $\alpha$-Tsallis SoftArgmax with $\alpha$-Tsallis regularization (right, $\alpha=0.1$). The value of the classical coupling Entropy--Softmax is much flatter than for our proposed coupling Tsallis--Tsallis. Using the latter removes flat areas that are far from the solution, allowing policy gradient methods to escape the gravitational pull.
  • Figure 2: Learning curves for Noisy CartPole (top row) and DeepSea (bottom row) under different choices of the Tsallis parameter $\alpha$. For Noisy CartPole, we report the standard unnoisedCartPole environment (a) and reward–noisy variants with increasing noise levels (b–d). For DeepSea, we consider grid sizes $L\in\{20,30,40,50\}$ (e–h). Each curve corresponds to the best temperature and step-size for a given $\alpha$, and shaded regions indicate $\pm$ one standard error over $25$ seeds. On Noisy CartPole, values $\alpha<1$ consistently improve performance over the PPO baseline in the standard and low-noise settings, with the gap increasing as the reward noise grows. On DeepSea, the improvement over the PPO baseline becomes more pronounced with increasing $L$, where $\alpha=0.7$ achieves the highest returns and the fastest learning.
  • Figure 3: Regularized value landscapes for a one-state, two-action MDP (rewards $0,1$) for different coupling between parameterizations and regularizations.
  • Figure 4: Average return as a function of training iterations on NChain (top row: sizes 10, 15, 20) and DeepSea (bottom row: sizes 10, 15). For \ref{['algo:PG']}, we report the best configuration for each divergence parameter $\alpha$; for all other methods, we report the best-performing configuration over their respective hyperparameters. Curves show the mean performance over $15$ independent seeds, with shaded regions indicating one standard deviation.

Theorems & Definitions (89)

  • Remark 3.1
  • Lemma 3.2
  • Theorem 3.3
  • Theorem 3.4
  • Theorem 3.5
  • Corollary 3.6
  • Remark 4.1: Connection with (Lazy) Mirror Descent.
  • Lemma 4.2
  • Theorem 4.3
  • Corollary 4.4
  • ...and 79 more