Why Policy Gradient Algorithms Work for Undiscounted Total-Reward MDPs

TL;DR

This work develops policy-gradient theory for undiscounted total-reward infinite-horizon MDPs by introducing a recurrent-transient decomposition and a transient visitation measure to overcome discontinuities that arise when gamma equals 1. It proves that state classifications are invariant on the interior policy set Pi_+ and derives a transient policy-gradient formula, enabling nonasymptotic convergence results for both projected policy gradient and natural policy gradient in the tabular setting. The results show sublinear convergence for projected policy gradient and either sublinear or linear convergence for natural policy gradient depending on step-size regimes, with experiments on Frozenlake and CliffWalk validating the theory and illustrating differences from pathological cases. Overall, the framework extends policy-gradient guarantees to undiscounted horizons, with implications for RL in large language model contexts and RLHF-like objectives.

Abstract

The classical policy gradient method is the theoretical and conceptual foundation of modern policy-based reinforcement learning (RL) algorithms. Most rigorous analyses of such methods, particularly those establishing convergence guarantees, assume a discount factor $γ< 1$. In contrast, however, a recent line of work on policy-based RL for large language models uses the undiscounted total-reward setting with $γ= 1$, rendering much of the existing theory inapplicable. In this paper, we provide analyses of the policy gradient method for undiscounted expected total-reward infinite-horizon MDPs based on two key insights: (i) the classification of the MDP states into recurrent and transient states is invariant over the set of policies that assign strictly positive probability to every action (as is typical in deep RL models employing a softmax output layer) and (ii) the classical state visitation measure (which may be ill-defined when $γ= 1$) can be replaced with a new object that we call the transient visitation measure.

Figures (3)

• Figure 1: Pathological MDP: The value function is discontinuous at the optimal policy, and the optimal action-value function does not specify the optimal policy
• Figure 2: Comparison of projected policy gradient (PPG) and natural policy gradient (NPG) algorithms in (left) Frozenlake and (right) Cliffwalk. The limit of the projected policy gradient algorithm gets closer to the optimum as $\alpha>0$ gets smaller.
• Figure 3: Comparison of projected policy gradient (PPG) and natural policy gradient (NPG) algorithms in Pathological MDP. Due to a discontinuity at optimal policy, $V^\star_\mu-V_\mu^{\pi_{k}} \ge V^\star_\mu - V^\star_{+,\mu} >0$.

Theorems & Definitions (46)

• Definition
• Proposition 1
• Lemma 1
• Lemma 2
• Lemma 3
• Lemma 4
• Lemma 5: Transient performance difference lemma
• Theorem 1: Transient policy gradient
• Lemma 6: Smoothness of value function
• Theorem 2