Landscape of Policy Optimization for Finite Horizon MDPs with General State and Action
Xin Chen, Yifan Hu, Minda Zhao
TL;DR
This work develops a unified Kurdyka-Łojasiewicz (KŁ) framework to establish global convergence for policy gradient methods solving finite-horizon MDPs with general state and action spaces. By verifying bounded gradients, the KŁ property of expected Q-values, and a sequential decomposition inequality, the paper shows that both exact and stochastic policy gradient methods converge to globally optimal policies with rates that scale polynomially in the planning horizon $T$; stochastic methods achieve an $\tilde{\mathcal{O}}(\epsilon^{-1})$ sample complexity. The framework applies to entropy-regularized tabular MDPs, linear quadratic regulators (LQR) with affine policies, multi-period inventory systems with Markov-modulated demands, and stochastic cash balance problems, providing the first known sample complexities for some of these settings. Across these applications, the KŁ constant depends polynomially on $T$, yielding non-asymptotic convergence guarantees despite nonconvex policy landscapes. The results offer a principled path to data-driven policy optimization in operations research contexts, with implications for scalable, provably optimal decision-making in dynamic systems.
Abstract
Policy gradient methods are widely used in reinforcement learning. Yet, the nonconvexity of policy optimization imposes 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 develop a framework that provides a set of easily verifiable assumptions to ensure the Kurdyka-Lojasiewicz (KL) condition of the policy optimization. Leveraging the KL condition, policy gradient methods converge to the globally optimal policy with a non-asymptomatic rate despite nonconvexity. Our results find applications in various control and operations models, including entropy-regularized tabular MDPs, Linear Quadratic Regulator (LQR) problems, stochastic inventory models, and stochastic cash balance problems, for which we show an $ε$-optimal policy can be obtained using a sample size in $\tilde{\mathcal{O}}(ε^{-1})$ and polynomial in terms of the planning horizon by stochastic policy gradient methods. Our result establishes the first sample complexity for multi-period inventory systems with Markov-modulated demands and stochastic cash balance problems in the literature.