Optimal Regret for Policy Optimization in Contextual Bandits

We present the first high-probability optimal regret bound for a policy optimization technique applied to the problem of stochastic contextual multi-armed bandit (CMAB) with general offline function approximation. Our algorithm is both efficient and achieves an optimal regret bound of $\widetilde{O}(\sqrt{ K|\mathcal{A}|\log|\mathcal{F}|})$, where $K$ is the number of rounds, $\mathcal{A}$ is the set of arms, and $\mathcal{F}$ is the function class used to approximate the losses. Our results bridge the gap between theory and practice, demonstrating that the widely used policy optimization methods for the contextual bandit problem can achieve a rigorously-proved optimal regret bound. We support our theoretical results with an empirical evaluation of our algorithm.

Contextual bandits model decision-making problems where an AI system must repeatedly choose among several actions based on the current context, such as recommending content to different users or selecting treatments for different patients. The key challenge is learning a decision strategy that quickly adapts to changing contexts while balancing exploration and performance. Although policy optimization methods are widely used in practice because they are simple and scalable, their theoretical guarantees have remained limited. In this work, we develop the first policy optimization algorithm for contextual bandits with general function approximation that achieves the best possible learning guarantees with high probability. Policy optimization methods improve decision-making directly by updating the strategy used to select actions, making them computationally attractive for large-scale applications. Our results show that these practical methods can also be theoretically optimal, matching the best achievable performance while remaining efficient to run. We further support our theory with experiments demonstrating the effectiveness of our approach in practice. Overall, this work helps close the gap between scalable algorithms used in real systems and the strong guarantees desired in reinforcement learning theory.