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

We introduce $\mathtt{\text{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 $\tilde{O}(H^4\sqrt{T\mathrm{\mid }S\mathrm{\mid }\mathrm{\mid }A\mathrm{\mid }\log (\mathrm{\mid }\mathcal{F}\mathrm{\mid }\mathrm{\mid }\mathcal{P}\mathrm{\mid })}),$ 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 $\mathrm{\mid }S\mathrm{\mid }$ and $\mathrm{\mid }A\mathrm{\mid }$, 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.