Beyond Stationarity: Convergence Analysis of Stochastic Softmax Policy Gradient Methods

TL;DR

This paper introduces a dynamic policy gradient method for finite-horizon MDPs that trains parameters backwards in time, improving convergence analysis and exploiting problem structure better than traditional methods.

Contribution

It proposes a novel dynamic policy gradient approach that incorporates backward training in finite-horizon MDPs and provides convergence analysis for this method.

Findings

Dynamic policy gradient outperforms standard methods in convergence speed.

The approach better exploits the structure of finite-time horizon problems.

Convergence bounds are improved using dynamic training.

Abstract

Markov Decision Processes (MDPs) are a formal framework for modeling and solving sequential decision-making problems. In finite-time horizons such problems are relevant for instance for optimal stopping or specific supply chain problems, but also in the training of large language models. In contrast to infinite horizon MDPs optimal policies are not stationary, policies must be learned for every single epoch. In practice all parameters are often trained simultaneously, ignoring the inherent structure suggested by dynamic programming. This paper introduces a combination of dynamic programming and policy gradient called dynamic policy gradient, where the parameters are trained backwards in time. For the tabular softmax parametrisation we carry out the convergence analysis for simultaneous and dynamic policy gradient towards global optima, both in the exact and sampled gradient settings without regularisation. It turns out that the use of dynamic policy gradient training much better exploits the structure of finite- time problems which is reflected in improved convergence bounds.