Fast Convergence of Softmax Policy Mirror Ascent
This work addresses the need for efficient and scalable policy optimization algorithms in reinforcement learning, offering incremental improvements over existing methods like NPG and MDPO.
The paper tackles the problem of policy optimization in reinforcement learning by refining a dual-space mirror ascent method (SPMA) to remove action normalization and prove faster convergence than existing methods. It demonstrates that SPMA achieves linear convergence in tabular MDPs and similar or better performance on benchmarks like MuJoCo and Atari compared to MDPO, PPO, and TRPO.
Natural policy gradient (NPG) is a common policy optimization algorithm and can be viewed as mirror ascent in the space of probabilities. Recently, Vaswani et al. [2021] introduced a policy gradient method that corresponds to mirror ascent in the dual space of logits. We refine this algorithm, removing its need for a normalization across actions and analyze the resulting method (referred to as SPMA). For tabular MDPs, we prove that SPMA with a constant step-size matches the linear convergence of NPG and achieves a faster convergence than constant step-size (accelerated) softmax policy gradient. To handle large state-action spaces, we extend SPMA to use a log-linear policy parameterization. Unlike that for NPG, generalizing SPMA to the linear function approximation (FA) setting does not require compatible function approximation. Unlike MDPO, a practical generalization of NPG, SPMA with linear FA only requires solving convex softmax classification problems. We prove that SPMA achieves linear convergence to the neighbourhood of the optimal value function. We extend SPMA to handle non-linear FA and evaluate its empirical performance on the MuJoCo and Atari benchmarks. Our results demonstrate that SPMA consistently achieves similar or better performance compared to MDPO, PPO and TRPO.