Beyond Softmax and Entropy: Convergence Rates of Policy Gradients with f-SoftArgmax Parameterization & Coupled Regularization

TL;DR

This paper introduces a new policy parameterization based on f-softargmax and coupled regularization, providing the first explicit convergence guarantees for policy gradient methods without preconditioning.

Contribution

It proposes an alternative to softmax parameterization using f-softargmax with regularization, enabling polynomial sample complexity and convergence guarantees.

Findings

f-PG with Tsallis divergences achieves polynomial sample complexity

Regularized f-softargmax improves optimization landscape and convergence

First non-asymptotic last-iterate convergence guarantees for policy gradients

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.