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

Orin Levy
ICML (2026)

Abstract

In this paper we introduce \emph{OPO-CMDP}, the first policy optimization algorithm for stochastic Contextual Markov Decision Process (CMDPs) under general offline function approximation. We establish a high probability regret bound of
$\widetilde{O}\left(H^4\sqrt{T|S||A|\log(|\mathcal{F}||\mathcal{P}|)}\right),$
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 function classes used to approximate the losses and dynamics, respectively. This result improves the dependence on $|S|$ and $|A|$ compared to the previously known state-of-the-art bound of \citet*{DBLP:conf/nips/QianHS24}. Our analysis introduces sophisticated confidence bounds for stochastic policies that eliminate restrictive assumptions required in prior work. Our results demonstrate that simple policy optimization over optimistic model approximations can achieve better, near-optimal regret bound for CMDPs.
×