Landscape of Policy Optimization for Finite Horizon MDPs with General State and Action

Abstract

Policy gradient methods are widely used in reinforcement learning. Yet, the nonconvexity of policy optimization poses significant challenges in understanding the global convergence of policy gradient methods. For a class of finite-horizon Markov Decision Processes (MDPs) with general state and action spaces, we identify a set of structural properties to establish a benign nonconvex landscape, the Polyak-{\L}ojasiewicz-Kurdyka (P{\L}K) condition of the policy optimization. Leveraging the P{\L}K condition, policy gradient methods converge to the globally optimal policy with a non-asymptotic rate despite nonconvexity. Our results apply to various control and operations models, including entropy-regularized tabular MDPs, Linear Quadratic Regulator problems, and both stochastic inventory models and stochastic cash balance problems with strongly convex costs. In these models, stochastic policy gradient methods obtain an $\epsilon$-optimal policy using a sample size of $\tilde{\mathcal{O}}(\epsilon^{-1})$ and polynomial in terms of the planning horizon. To the best of our knowledge, we provide the first sample-complexity guarantees for multi-period inventory systems with Markov-modulated demand and for stochastic cash balance problems. We complement the theory with numerical experiments showing that policy gradient methods outperform several benchmark algorithms from the literature across these operations models.