SPEED: Experimental Design for Policy Evaluation in Linear Heteroscedastic Bandits
Subhojyoti Mukherjee, Qiaomin Xie, Josiah Hanna, Robert Nowak
TL;DR
The paper tackles policy evaluation in linear bandits with heteroscedastic noise by formulating a PE-Optimal design that minimizes the MSE of the target policy value $v(\pi)$. It introduces SPEED, an agnostic algorithm that learns the covariance $\mathbf{\Sigma}_*$ during an exploration phase and then follows the PE-Optimal design to collect data, achieving a regret of $O_{\kappa^2,H^2_U}\big( \frac{d^3 \sigma^2_{\max} \log n}{\sigma^2_{\min} n^{3/2}} \big)$ and matching a lower bound up to a factor of $d$. The design leverages a weighted least squares framework with $\tilde{\mathbf{x}}(a) = \mathbf{x}(a)/\sigma(a)$ and a novel convex optimization over designs $\mathbf{b}$. Empirically, SPEED attains policy evaluation MSE close to an oracle that knows $\mathbf{\Sigma}_*$ and outperforms on-policy sampling across several datasets. The work advances practical data-collection strategies for reliable policy evaluation in structured, variance-varying environments with implications for safer and more efficient online deployment.
Abstract
In this paper, we study the problem of optimal data collection for policy evaluation in linear bandits. In policy evaluation, we are given a target policy and asked to estimate the expected reward it will obtain when executed in a multi-armed bandit environment. Our work is the first work that focuses on such optimal data collection strategy for policy evaluation involving heteroscedastic reward noise in the linear bandit setting. We first formulate an optimal design for weighted least squares estimates in the heteroscedastic linear bandit setting that reduces the MSE of the value of the target policy. We then use this formulation to derive the optimal allocation of samples per action during data collection. We then introduce a novel algorithm SPEED (Structured Policy Evaluation Experimental Design) that tracks the optimal design and derive its regret with respect to the optimal design. Finally, we empirically validate that SPEED leads to policy evaluation with mean squared error comparable to the oracle strategy and significantly lower than simply running the target policy.