Near-Optimal Regret for Policy Optimization in Contextual MDPs with General Offline Function Approximation
Orin Levy, Aviv Rosenberg, Alon Peled-Cohen, Yishay Mansour
Abstract
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.
BibTeX
@inproceedings{
levy2026nearoptimal,
title={Near-Optimal Regret for Policy Optimization in Contextual {MDP}s with General Offline Function Approximation},
author={Orin Levy and Aviv Rosenberg and Alon Cohen and Yishay Mansour},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=5faxBOZAmW}
}