Sample Complexity of Neural Policy Mirror Descent for Policy Optimization on Low-Dimensional Manifolds
Zhenghao Xu, Xiang Ji, Minshuo Chen, Mengdi Wang, Tuo Zhao
TL;DR
The work addresses why deep policy gradient methods using CNNs can tackle high-dimensional RL challenges by exploiting low-dimensional state-space structure. It introduces Neural Policy Mirror Descent (NPMD) with an actor-critic CNN architecture and proves that both iteration and sample complexities depend on the intrinsic dimension $d$ of the state manifold rather than the ambient space dimension $D$, achieving a bound of $ ilde{O}( ext{ε}^{-d/α-2})$ under Lipschitz MDP assumptions. The paper develops a robust CNN approximation theory on manifolds, derives explicit CNN-size requirements, and couples them with probabilistic arguments to obtain tight sample-complexity bounds for both critic and actor updates. Numerical experiments on image-based CartPole demonstrate that the method remains effective with increasing ambient dimension, aligning with the theory that intrinsic geometry governs learning efficiency. Overall, the results provide a principled explanation for the empirical success of deep policy gradient methods in high-dimensional tasks by leveraging low-dimensional state-space structures.
Abstract
Policy gradient methods equipped with deep neural networks have achieved great success in solving high-dimensional reinforcement learning (RL) problems. However, current analyses cannot explain why they are resistant to the curse of dimensionality. In this work, we study the sample complexity of the neural policy mirror descent (NPMD) algorithm with deep convolutional neural networks (CNN). Motivated by the empirical observation that many high-dimensional environments have state spaces possessing low-dimensional structures, such as those taking images as states, we consider the state space to be a $d$-dimensional manifold embedded in the $D$-dimensional Euclidean space with intrinsic dimension $d\ll D$. We show that in each iteration of NPMD, both the value function and the policy can be well approximated by CNNs. The approximation errors are controlled by the size of the networks, and the smoothness of the previous networks can be inherited. As a result, by properly choosing the network size and hyperparameters, NPMD can find an $ε$-optimal policy with $\widetilde{O}(ε^{-\frac{d}α-2})$ samples in expectation, where $α\in(0,1]$ indicates the smoothness of environment. Compared to previous work, our result exhibits that NPMD can leverage the low-dimensional structure of state space to escape from the curse of dimensionality, explaining the efficacy of deep policy gradient algorithms.