# Best Arm Identification with Contextual Information under a Small Gap

arXiv: https://arxiv.org/abs/2209.07330 · HTML: https://mlsanity.com/paper/2209.07330/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 105 references that are in the map.
- First posted on arXiv 15 Sep 2022.

## Builds on (most cited first, 30 of 34)

- [On the Complexity of Best Arm Identification in Multi-Armed Bandit Models](https://mlsanity.com/paper/1407.4443/) (2014)
- [Optimal Best Arm Identification with Fixed Confidence](https://mlsanity.com/paper/1602.04589/) (2016)
- [Time-uniform, nonparametric, nonasymptotic confidence sequences](https://mlsanity.com/paper/1810.08240/) (2018)
- [lil' UCB : An Optimal Exploration Algorithm for Multi-Armed Bandits](https://mlsanity.com/paper/1312.7308/) (2013)
- [Simple Bayesian Algorithms for Best Arm Identification](https://mlsanity.com/paper/1602.08448/) (2016)
- [Confidence Intervals for Policy Evaluation in Adaptive Experiments](https://mlsanity.com/paper/1911.02768/) (2019)
- [Tight (Lower) Bounds for the Fixed Budget Best Arm Identification Bandit Problem](https://mlsanity.com/paper/1605.09004/) (2016)
- [Mixture Martingales Revisited with Applications to Sequential Tests and Confidence Intervals](https://mlsanity.com/paper/1811.11419/) (2018)
- [Non-Asymptotic Pure Exploration by Solving Games](https://mlsanity.com/paper/1906.10431/) (2019)
- [Improving the Expected Improvement Algorithm](https://mlsanity.com/paper/1705.10033/) (2017)
- [An Empirical Process Approach to the Union Bound: Practical Algorithms for Combinatorial and Linear Bandits](https://mlsanity.com/paper/2006.11685/) (2020)
- [Off-Policy Evaluation via Adaptive Weighting with Data from Contextual Bandits](https://mlsanity.com/paper/2106.02029/) (2021)
- [Fixed-Confidence Guarantees for Bayesian Best-Arm Identification](https://mlsanity.com/paper/1910.10945/) (2019)
- [The Influence Function of Semiparametric Estimators](https://mlsanity.com/paper/1508.01378/) (2015)
- [Top Two Algorithms Revisited](https://mlsanity.com/paper/2206.05979/) (2022)
- [Post-Contextual-Bandit Inference](https://mlsanity.com/paper/2106.00418/) (2021)
- [Online Multi-Armed Bandits with Adaptive Inference](https://mlsanity.com/paper/2102.13202/) (2021)
- [Doubly robust Thompson sampling for linear payoffs](https://mlsanity.com/paper/2102.01229/) (2021)
- [Sequential Nonparametric Testing with the Law of the Iterated Logarithm](https://mlsanity.com/paper/1506.03486/) (2015)
- [Minimax Optimal Algorithms for Fixed-Budget Best Arm Identification](https://mlsanity.com/paper/2206.04646/) (2022)
- [Policy Learning with Observational Data](https://mlsanity.com/paper/1702.02896/) (2017)
- [Nonparametric Stochastic Contextual Bandits](https://mlsanity.com/paper/1801.01750/) (2018)
- [A/B/n Testing with Control in the Presence of Subpopulations](https://mlsanity.com/paper/2110.15573/) (2021)
- [Simple Regret Minimization for Contextual Bandits](https://mlsanity.com/paper/1810.07371/) (2018)
- [Policy Choice and Best Arm Identification: Asymptotic Analysis of Exploration Sampling](https://mlsanity.com/paper/2109.08229/) (2021)
- [On the Complexity of A/B Testing](https://mlsanity.com/paper/1405.3224/) (2014)
- [Stratification Trees for Adaptive Randomization in Randomized Controlled Trials](https://mlsanity.com/paper/1806.05127/) (2018)
- [Sample complexity of partition identification using multi-armed bandits](https://mlsanity.com/paper/1811.05654/) (2018)
- [Diffusion Approximations for Thompson Sampling in the Small Gap Regime](https://mlsanity.com/paper/2105.09232/) (2021)
- [Efficient Online Estimation of Causal Effects by Deciding What to Observe](https://mlsanity.com/paper/2108.09265/) (2021)

## Built on by (most cited first, 3 of 3)

- [Towards Regulatory-Confirmed Adaptive Clinical Trials: Machine Learning Opportunities and Solutions](https://mlsanity.com/paper/2503.09226/) (2025)
- [Conformal bandits: bringing statistical validity and reward efficiency under weak arm separability](https://mlsanity.com/paper/2512.09850/) (2025)
- [Active Context Selection Improves Simple Regret in Contextual Bandits](https://mlsanity.com/paper/2605.20040/) (2026)

## Abstract (arXiv)

We study the best-arm identification (BAI) problem with a fixed budget and contextual (covariate) information. In each round of an adaptive experiment, after observing contextual information, we choose a treatment arm using past observations and current context. Our goal is to identify the best treatment arm, which is a treatment arm with the maximal expected reward marginalized over the contextual distribution, with a minimal probability of misidentification. In this study, we consider a class of nonparametric bandit models that converge to location-shift models when the gaps go to zero. First, we derive lower bounds of the misidentification probability for a certain class of strategies and bandit models (probabilistic models of potential outcomes) under a small-gap regime. A small-gap regime is a situation where gaps of the expected rewards between the best and suboptimal treatment arms go to zero, which corresponds to one of the worst cases in identifying the best treatment arm. We then develop the ``Random Sampling (RS)-Augmented Inverse Probability weighting (AIPW) strategy,'' which is asymptotically optimal in the sense that the probability of misidentification under the strategy matches the lower bound when the budget goes to infinity in the small-gap regime. The RS-AIPW strategy consists of the RS rule tracking a target sample allocation ratio and the recommendation rule using the AIPW estimator.