# Multi-Armed Bandits With Machine Learning-Generated Surrogate Rewards

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

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 36 references that are in the map.
- First posted on arXiv 20 Jun 2025.

## Builds on (most cited first, 15 of 15)

- [A Tutorial on Thompson Sampling](https://mlsanity.com/paper/1707.02038/) (2017)
- [Prediction-Powered Inference](https://mlsanity.com/paper/2301.09633/) (2023)
- [PPI++: Efficient Prediction-Powered Inference](https://mlsanity.com/paper/2311.01453/) (2023)
- [Multi-armed Bandit Models for the Optimal Design of Clinical Trials: Benefits and Challenges](https://mlsanity.com/paper/1507.08025/) (2015)
- [Active Statistical Inference](https://mlsanity.com/paper/2403.03208/) (2024)
- [Semi-supervised Inference: General Theory and Estimation of Means](https://mlsanity.com/paper/1606.07268/) (2016)
- [Predictions as Surrogates: Revisiting Surrogate Outcomes in the Age of AI](https://mlsanity.com/paper/2501.09731/) (2025)
- [On Multi-Armed Bandit Designs for Dose-Finding Clinical Trials](https://mlsanity.com/paper/1903.07082/) (2019)
- [Warm-starting Contextual Bandits: Robustly Combining Supervised and Bandit Feedback](https://mlsanity.com/paper/1901.00301/) (2019)
- [Prediction De-Correlated Inference: A safe approach for post-prediction inference](https://mlsanity.com/paper/2312.06478/) (2023)
- [Artificial Replay: A Meta-Algorithm for Harnessing Historical Data in Bandits](https://mlsanity.com/paper/2210.00025/) (2022)
- [Long-Term Value of Exploration: Measurements, Findings and Algorithms](https://mlsanity.com/paper/2305.07764/) (2023)
- [Online Bandits with (Biased) Offline Data: Adaptive Learning under Distribution Mismatch](https://mlsanity.com/paper/2405.02594/) (2024)
- [Stochastic Multi-Armed Bandits with Control Variates](https://mlsanity.com/paper/2105.03962/) (2021)
- [Another look at statistical inference with machine learning-imputed data](https://mlsanity.com/paper/2411.19908/) (2024)

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

- [OptiRepair: Closed-Loop Diagnosis and Repair of Supply Chain Optimization Models with LLM Agents](https://mlsanity.com/paper/2602.19439/) (2026)
- [Contextual Online Pricing with (Biased) Offline Data](https://mlsanity.com/paper/2507.02762/) (2025)
- [Revisiting Active Sequential Prediction-Powered Mean Estimation](https://mlsanity.com/paper/2604.18569/) (2026)
- [Valid Best-Model Identification for LLM Evaluation via Low-Rank Factorization](https://mlsanity.com/paper/2605.10405/) (2026)
- [Correlation-Aware Contextual Bandits with Surrogate Rewards for LLM Routing](https://mlsanity.com/paper/2607.09015/) (2026)
- [When Should Humans Step In? Optimal Human Dispatching in AI-Assisted Decisions](https://mlsanity.com/paper/2603.13688/) (2026)
- [Best-Arm Identification with Generative Proxy](https://mlsanity.com/paper/2607.06879/) (2026)
- [Perturbing the Derivative: Wild Refitting for Model-Free Evaluation of Machine Learning Models under Bregman Losses](https://mlsanity.com/paper/2509.02476/) (2025)
- [Why Most Optimism Bandit Algorithms Have the Same Regret Analysis: A Simple Unifying Theorem](https://mlsanity.com/paper/2512.18409/) (2025)
- [Uncertainty-aware Causal Decision Making via Effect Bound Decomposition](https://mlsanity.com/paper/2601.22736/) (2026)
- [When May a Model Replace the Experiment? Audits, Licenses, and the Price of Trust in Surrogate-Driven Design](https://mlsanity.com/paper/2608.01378/) (2026)
- [Learning a Size-Weight Frontier for Synthetic-Augmented Inference](https://mlsanity.com/paper/2608.28576/) (2026)
- [Augmented Hypothesis Testing with Persona-Based LLM Simulations](https://mlsanity.com/paper/2609.24629/) (2026)
- [Gaussian Process Bandits for AI System Optimization with Selectively Verified Feedback](https://mlsanity.com/paper/2610.17760/) (2026)

## Abstract (arXiv)

Multi-armed bandit (MAB) is a widely adopted framework for sequential decision-making under uncertainty. Traditional bandit algorithms rely solely on online data, which tends to be scarce as it must be gathered during the online phase when the arms are actively pulled. However, in many practical settings, rich auxiliary data, such as covariates of past users, is available prior to deploying any arms. We introduce a new setting for MAB where pre-trained machine learning (ML) models are applied to convert side information and historical data into surrogate rewards. A prominent challenge of this setting is that the surrogate rewards may exhibit substantial bias, as true reward data is typically unavailable in the offline phase, forcing ML predictions to heavily rely on extrapolation. To address the issue, we propose the Machine Learning-Assisted Upper Confidence Bound (MLA-UCB) algorithm, which can be applied to any reward prediction model and any form of auxiliary data. When the predicted and true rewards are jointly Gaussian, it provably improves the cumulative regret, even in cases where the mean surrogate reward completely misaligns with the true mean rewards, and achieves the asymptotic optimality among a broad class of policies. Notably, our method requires no prior knowledge of the covariance matrix between true and surrogate rewards. We further extend the method to a batched reward MAB problem, where each arm pull yields a batch of observations and rewards may be non-Gaussian, and we derive computable confidence bounds and regret guarantees that improve upon classical UCB algorithms. Finally, extensive simulations with both Gaussian and ML-generated surrogates, together with real-world studies on language model selection and video recommendation, demonstrate consistent and often substantial regret reductions with moderate offline surrogate sample sizes and correlations.