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

TL;DR

This paper presents exttt{OPO-CMDP}, a novel policy optimization algorithm for stochastic CMDPs with offline function approximation, achieving near-optimal regret bounds and improving upon previous methods in terms of dependence on state and action spaces.

Contribution

Introduces exttt{OPO-CMDP}, the first policy optimization algorithm with regret bounds for CMDPs under general offline function approximation, with optimal dependence on state and action spaces.

Findings

Achieves high probability regret bound of ilde{O}(H^4√(T|S||A|log(|| |P|)))

First regret bound with optimal dependence on |S| and |A|

Demonstrates that optimistic policy optimization is computationally efficient and near-optimal for CMDPs.

Abstract

We introduce \texttt{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.