Near-Optimal Regret for Policy Optimization in Contextual MDPs with General Offline Function Approximation

We introduce OPO-CMDP, the first policy optimization algorithm for stochastic Contextual Markov Decision Process (CMDPs) under general offline function approximation. Our approach achieves a high probability regret bound of $\widetilde{O}(H^4\sqrt{T|S||A|\log(|\mathcal{F}||\mathcal{P}|)}),$ where $S$ and $A$ denote the state and action spaces, $H$ the horizon length, $T$ the number of episodes, and $\mathcal{F}, \mathcal{P}$ the finite function classes used to approximate the losses and dynamics, respectively. This is the first regret bound with optimal dependence on $|S|$ and $|A|$, directly improving the current state-of-the-art (Qian, Hu, and Simchi-Levi, 2024). These results demonstrate that optimistic policy optimization provides a natural, computationally superior and theoretically near-optimal path for solving CMDPs.

Contextual reinforcement learning studies how an AI system can learn to make good sequential decisions when each situation may have different characteristics or “contexts,” such as different users, environments, or tasks. A central challenge is learning effective decision policies efficiently while handling uncertainty and limited feedback. Existing approaches often require heavy computation or fail to scale well to large problems. In this work, we develop a new policy optimization method, called OPO-CMDP, for contextual reinforcement learning. Policy optimization methods improve decision-making directly by updating the strategy the agent uses to choose actions. Our algorithm combines this approach with principled exploration, allowing it to learn efficiently even in complex environments with large state and action spaces. We prove strong theoretical guarantees showing that our method learns near-optimally over time, while improving the best previous results in the literature. Beyond the theory, our work shows that policy optimization can be both computationally practical and mathematically reliable for contextual reinforcement learning, making it a promising direction for building scalable decision-making systems in real-world applications.