On the Global Optimality of Policy Gradient Methods in General Utility Reinforcement Learning
Anas Barakat, Souradip Chakraborty, Peihong Yu, Pratap Tokekar, Amrit Singh Bedi
TL;DR
This work addresses global optimality of policy gradient methods for reinforcement learning with general utilities (RLGU), encompassing frameworks beyond standard returns. It proves a gradient-domination property in the tabular setting under concave $F$, enabling global optimality guarantees for stationary points, and then proposes a scalable two-loop algorithm (PG-OMA) that estimates occupancy measures via maximum likelihood in a function-approximation class. The analysis delivers nonconcave and concave-utility guarantees, showing that with occupancy-measure estimation error bounded by $\epsilon_{MLE}$, one can achieve a last-iterate global optimality bound and a total sample complexity of $\tilde{O}(m \varepsilon^{-4})$, where $m$ is the density-feature dimension. Practically, this work enables scalable RLGU solutions by decoupling occupancy estimation from policy updates and leveraging MLE to control statistical complexity, thus widening the applicability of policy-gradient methods to large state-action spaces and general utilities.
Abstract
Reinforcement learning with general utilities (RLGU) offers a unifying framework to capture several problems beyond standard expected returns, including imitation learning, pure exploration, and safe RL. Despite recent fundamental advances in the theoretical analysis of policy gradient (PG) methods for standard RL and recent efforts in RLGU, the understanding of these PG algorithms and their scope of application in RLGU still remain limited. In this work, we establish global optimality guarantees of PG methods for RLGU in which the objective is a general concave utility function of the state-action occupancy measure. In the tabular setting, we provide global optimality results using a new proof technique building on recent theoretical developments on the convergence of PG methods for standard RL using gradient domination. Our proof technique opens avenues for analyzing policy parameterizations beyond the direct policy parameterization for RLGU. In addition, we provide global optimality results for large state-action space settings beyond prior work which has mostly focused on the tabular setting. In this large scale setting, we adapt PG methods by approximating occupancy measures within a function approximation class using maximum likelihood estimation. Our sample complexity only scales with the dimension induced by our approximation class instead of the size of the state-action space.