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  (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  leads to policy evaluation with mean squared error comparable to the oracle strategy and significantly lower than simply running the target policy.

§ INTRODUCTION

Bandit policy optimization has been widely applied in diverse applications such as web marketing Bottou et al., 2013, web search Li et al., 2011, and healthcare recommendations Zhou et al., 2017.

In practice, before widely deploying a learned policy, it is necessary to have an accurate estimation of its performance (i.e., expected reward).

To this effect, policy evaluation is often a critical step as it allows practitioners to determine if a learned policy truly represents improved task performance.

While off-policy evaluation (OPE) has been extensively studied as a potential solution [Dudík et al., 2014, Li et al., 2015, Swaminathan et al., 2017, Wang et al., 2017, Su et al., 2020, Kallus et al., 2021, Cai et al., 2021], in practice, some amount of limited, online evaluation is often required before more widescale deployment.

For instance, in web-marketing it is common to run an A/B test with a subset of all users before a potential new policy is deployed for all users [Kohavi and Longbotham, 2017].

When online policy evaluation is required, we desire methods that provide an accurate estimate of policy performance with a minimal amount of data collected.

The default choice for online policy evaluation is to simply run the target policy and average the resulting rewards.

However, this approach is sub-optimal when the space of actions is large or different actions have reward distributions with different variances.

In this paper, we formulate a new experimental design for allocating action samples so as to obtain minimal mean squared error policy evaluation.

Specifically, we consider optimal policy evaluation under the following linear heteroscedastic bandit model.

Let $\A$ be a set of actions.  Each $a \in A$ is associated with a vector $\bx(a)\in \R^d$. The expected reward of action $a$ is a linear function $\btheta_*^\top \bx(a)$, for some $\btheta_* \in \R^d$.  Often the variance of the reward is assumed to be the same for all actions, but in this paper, we depart from this assumption.  We suppose that the variance is governed by a quadratic function of the form $\bx(a)^T \bSigma_* \bx(a)$, for some symmetric positive definite matrix $\bSigma_*\in\R^{d\times d}$.  This assumption allows us to capture problems in which both the mean reward and the variance may depend on the action taken, but both vary smoothly in $\bx(a)$.

We briefly contrast our studied setting with other work.

Recently, there has been an increasing focus on theoretically optimal data collection for bandit policy evaluation.

In policy evaluation, the usual metric of algorithm performance is regret with respect to the mean squared error of an oracle algorithm that has knowledge of the variances of different reward distributions (i.e., knows $\Sigma^*$).
In recent times there has been an increasing focus on studying data collection for policy evaluation in bandit settings [Zhu and Kveton, 2021, Zhu and Kveton, 2022, Wan et al., 2022] and there has been some theoretical progress [Chaudhuri et al., 2017, Fontaine et al., 2021].

Several works [Antos et al., 2008, Carpentier and Munos, 2012, Carpentier et al., 2015, Fontaine et al., 2021] have shown that when $d = A$ a regret of $\widetilde{O}(A n^{-3/2})$ is possible where $n$ is the total budget of actions that can be tried and $\widetilde{O}$ hides logarithmic factors.

These works have also shown that simply running the target policy to take actions results in a slower decrease of regret at the rate of $\widetilde{O}(A n^{-1})$.

The work of Zhu and Kveton, 2022, Wan et al., 2022 studies the same setting under safety constraints and provides asymptotic error bounds.

However, none of the above works provides a finite-time regret guarantee for data collection for policy evaluation in the heteroscedastic linear bandit setting.

The closest works to our own [Antos et al., 2008, Carpentier and Munos, 2012, Carpentier et al., 2015, Fontaine et al., 2021] either consider unstructured settings or assume that $d = A$.

As many real-world bandit applications have $d \ll A$, a natural question arises as to how to build an algorithm for policy evaluation in the heteroscedastic linear bandit setting with unknown $\btheta_*$ and $\bSigma_*$ that can leverage the structure.

Further, we want the regret of such an algorithm to decrease at a faster rate than $\widetilde{O}(n^{-1})$ (the on-policy regret rate) and to scale with the dimension $d$ instead of actions as $A\gg d$. Note that the regret should scale at least by $d^2$ because the learner needs to probe in $d^2$ dimensions to estimate $\bSigma_*\in\R^{d\times d}$ [Wainwright, 2019]. Thus, the goal of our work is to answer the question:

Can we design an algorithm to collect data  for policy evaluation that adapts to the variance of each action, and its regret degrades at a faster rate than $\widetilde{O}(d^2n^{-1})$?

In this paper, we answer this question affirmatively.

We note that heteroscedasticity is also studied for policy improvement setup [Kirschner and Krause, 2018, Zhou and Gu, 2022, Zhou et al., 2021, Zhang et al., 2021, Zhao et al., 2022]. In these prior works the reward variances are time-dependent as opposed to the quadratic structure studied in this paper.

Note that policy improvement requires a different approach than policy evaluation. These works build tight confidence sets around the unknown model parameter $\btheta_*$ by employing weighted ridge regression involving an estimated upper bound to the time-dependent variances. However, in our setting, the variances of each action share the unknown low dimensional co-variance matrix $\bSigma_*$. Hence we deviate from these approaches and employ an alternating OLS-WLS estimation to learn the underlying parameter $\bSigma_*$. We discuss more related works and motivations in <Ref>.

We make the following novel contributions to the growing literature on online policy evaluation:

1. We are the first to formulate the policy evaluation problem for heteroscedastic linear bandit setting where the variance of each action $\bx(a)\in\R^d$ depends on the lower dimensional co-variance matrix $\bSigma_*\in\R^{d\times d}$ such that variance $\sigma^2(a) = \bx(a)^\top\bSigma_*\bx(a)$. This is a more general heteroscedastic linear bandit setting than studied in Chaudhuri et al., 2017, Kirschner and Krause, 2018, Fontaine et al., 2021, and different than time-dependent variance model of Zhang et al., 2021, Zhao et al., 2022.
2. We characterize our loss in this setting and show that the optimal design, denoted as Policy Evaluation (PE) Optimal  design to minimize this loss is different than  A-, D-, E-, G-, T-, V-optimality [Pukelsheim, 2006]. We establish several key properties of this novel design and discuss how we can solve this optimization problem efficiently.
3. Finally we propose the agnostic algorithm, ,  that tracks this optimal design and analyze

its MSE. We then bound the regret of  compared to an oracle strategy that follows the optimal design with the knowledge of $\bSigma_*$. We show that the regret scales as $O(\tfrac{d^3\log(n)}{n^{3/2}})$ which is an improvement over the regret for the stochastic non-structured bandit setting which scales as $O(\tfrac{A\log(n)}{n^{3/2}})$ [Carpentier and Munos, 2011, Carpentier and Munos, 2012, Carpentier et al., 2015, Fontaine et al., 2021]. Hence, we answer positively to our main query. We also prove the first lower bound for this setting that scales as $\Omega(\tfrac{d^2\log(n)}{n^{3/2}})$. Finally, we complement our theoretical findings with experiments on real-life data sets.

§ PRELIMINARIES

We define $[m]\coloneqq [1,2,\ldots,m]$. The setting consists of $A$ actions, indexed by $a\in[A]$, and consists of features $\bx(a)\in\R^d$ such that the dimension $d^2 \ll A$.

Denote by $\Delta(\A)$ the probability simplex over the action space $\A$ and a policy $\pi\in\Delta(\A)$ as a mapping $\pi: a \rightarrow [0,1]$ such that $\sum_a \pi(a) = 1$.

We denote the total available budget as $n$.

We study the linear bandit setting where the expected reward for each action is assumed to be a linear function [Mason et al., 2021, Jamieson and Jain, 2022].

Specifically, at each round $t\in[n]$, the selected action $a_t$ is associated with a feature vector $\bx(a_t) \in \R^d$, and the rewards satisfy:

$R_{t}(a_t) =  \bx(a_t)^{\top}\btheta_* + \eta$,

where $\btheta_*\in\R^d$ is the unknown reward parameter, and $\eta$ is zero-mean noise with variance $\sigma^2(a)$.

We assume that the variance $\sigma^2(a)$ has a lower-dimensional structure such that $\sigma^2(a) = \bv(a)^{\top}\bSigma_*\bv(a)$ where $\bSigma_* \in \R^{d\times d}$ is an unknown variance parameter and we further assume that $\eta$ is bounded between $-B$ to $B$. Observe that the variance depends on the action features, which is called the heteroscedastic noise model [Greene, 2002, Chaudhuri et al., 2017] and differs from the unknown time-dependent variance model of Zhang et al., 2021, Zhao et al., 2022.

Moreover, Chaudhuri et al., 2017 only consider the special case of our setting when $\bSigma_*$ is identity, such that $\sigma^2(a) = \bx(a)^{\top}\bx(a)$.

We also assume that the norms of the features are bounded such that $H^2_L \leq \|\bx(a)\|^2\leq H^2_U$ for all $a\in\A$.
In our heteroscedastic linear bandit setting selecting any action gives information about $\btheta_*$ and also gives information about the noise covariance matrix $\bSigma_*$.

The value of a policy $\pi$ is defined as $v(\pi) \coloneqq \E[R_t]$ where the expectation is taken over $a_{t}\!\sim\!\pi,R_{t}\!\sim\! R(a_{t})$.

Finally, recall that in the policy evaluation problem, we are given a fixed, target policy $\pi$ and asked to estimate $v(\pi)$.
Estimating $v(\pi)$ requires a dataset of actions and their associated rewards, $\D \coloneqq \{(a_1, r_1,...,a_{n},r_{n})\}$, which is collected by executing some policy.

We refer to the policy that collects $\D$ as the behavior policy, denoted by $\bb\in\triangle(\A)$.
We then define the value estimate of a policy $\pi$ as $Y_n$, where $n$ is the sample budget. The exact nature of the value estimate for the linear bandit setting will be made clear in <Ref>. Our goal is to choose a behavior policy that minimizes the mean squared error (MSE) defined as

$\E_{\mathcal{D}}[(Y_n - v(\pi))^{2}]$,

where the expectation is over the collected data set $\D$. In this paper, we use the terms MSE and loss interchangeably.

We  now

state the assumption on the boundedness on the variance of each action $a\in[A]$.

Let the  singular value decomposition of $\bSigma_*$ be $\bU \bD \bP^{\top}$ with orthogonal matrices $\bU, \bP^{\top}$ and $\bD=\operatorname{diag}\left(\lambda_{1}, \ldots, \lambda_{d}\right)$ where $\lambda_{i}$ denotes a singular value. It follows that $\sigma^2_{\min} \leq \sigma^2(a)\leq \sigma^2_{\max}$ where $\sigma^2_{\min} = \min_{i}|\lambda_{i}|H_L^2$ and $\sigma^2_{\max} = \max_{i}|\lambda_{i}|H^2_U$ (see <Ref>).

We assume that $\bSigma_*$ has its minimum and maximum eigenvalues bounded such that for every action $a\in[A]$ the following holds $\sigma^2_{\min} \leq \sigma^2(a)\leq \sigma^2_{\max}$.

Our work is most closely related to existing work on data collection for policy evaluation.

Perhaps the most natural choice of behavior policy is to simply run the target policy, i.e., on-policy data collection [Sutton and Barto, 2018].

The works in adaptive importance sampling for bandits [Oosterhuis and de Rijke, 2020, Tucker and Joachims, 2022] and MDPs [Hanna et al., 2017, Ciosek and Whiteson, 2017, Bouchard et al., 2016] have shown how to lower the variance of Monte Carlo estimation through the choice of behavior policy.

In contrast to these works, we consider estimating $v(\pi)$ by estimating the reward distributions rather than using Monte Carlo estimation.

Such certainty-equivalence estimators take advantage of the setting's structure and are thus typically of lower variance than Monte Carlo estimators [Sutton and Barto, 2018].

The work of Wan et al., 2022 studies a different estimator for reducing the variance of the importance sampling in constrained MDP setting whereas we study certainty equivalence estimator.

Another set of work has studied sample allocation for stratified Monte Carlo estimators – a problem that is formally equivalent to behavior policy selection for policy evaluation in the bandit setting with linearly independent arms [Antos et al., 2008, Carpentier et al., 2015].

This line of work was recently extended to tabular, tree-structured MDPs by Mukherjee et al., 2022.

In contrast, we consider the structured linear bandit setting which incorporates generalization across actions.

Our work is closely related to optimal experimental design and active learning literature.

We formulate determining the optimal behavior policy in the bandit setting as an optimal design problem.

In contrast to prior work, we introduce a new type of optimality that is tailored to the policy evaluation problem. We are also, to the best of our knowledge, the first to consider both heteroscedastic noise and weighted least squares estimators in formulating our design.

The heteroscedastic noise model and weighted least squares estimator have been considered by [Chaudhuri et al., 2017] in the active learning literature and in linear bandit setting by Kirschner and Krause, 2018 using information directed sampling.

In contrast to these works (and the active learning setting in general), we aim to minimize the weighted error $\sum_{a\in\bA} \pi(a) \bv(a)^\top (\btheta^* - \wtheta)^2$ whereas in the active learning setting the goal is to minimize $\|\btheta^* - \wtheta\|^2$ which results in A-optimal design [Fontaine et al., 2021, Pukelsheim, 2006]. Moreover the regret bounds in Fontaine et al., 2021 holds for $d=|\A|$.

[Riquelme et al., 2017] extends the results of [Carpentier and Munos, 2011] to a different linear regression setting than ours but under homoscedastic noise model

Data collection for policy evaluation is also related to the problem of exploration for policy learning in MDPs or best-arm identification in bandits.

In those contexts, the aim of exploration is to find the optimal policy and the exploration-exploitation trade-off describes the tension between reducing uncertainty and focusing on known promising actions.

In bandits, the exploration-exploitation trade-off is often navigated under the “Optimism in the Face of Uncertainty" principle using techniques such as UCB [Lai and Robbins, 1985, Auer et al., 2002, Abbasi-Yadkori et al., 2011] or Thompson Sampling [Thompson, 1933, Agrawal and Goyal, 2012].

In contrast to the standard exploration problem, we focus on evaluating a fixed policy.

Instead of balancing exploration and exploitation, a behavior policy for policy evaluation should take actions that reduce uncertainty about $v(\pi)$ with emphasis on actions that have high probability under $\pi$.
Also, note that heteroscedastic bandits have been studied from the perspective of policy improvement [Kirschner and Krause, 2018, Zhao et al., 2022] however, in this paper we focus on optimal data collection for policy evaluation.

§ OPTIMAL DESIGN FOR POLICY EVALUATION

In this section, we first derive an expression for policy evaluation error in terms of

the behavior sampling proportion $\bb\in\triangle(\A)$, target policy $\pi$, and action features $\bx(a)\in\R^d$.

We call this expression "optimal design" [Pukelsheim, 2006] as minimizing it results in minimizing the error for policy evaluation.

We then analyze the error incurred by an oracle

that has access to problem-dependent parameters.

Suppose we have a budget of $n$ samples to divide between the actions, and let $T_{n}(1), T_{n}(2), \ldots, T_{n}(A)$ be the number of samples allocated to actions $1,2, \ldots, A$ at the end of $n$ rounds.

In a linear bandit, we can define the value estimate of a target policy as

$Y_n \coloneqq \sum_a \bw(a)^\top \wtheta_n$
where, $\bw(a) \coloneqq \pi(a)\bv(a)$ is the expected features under the target policy, and $\wtheta_n$ is an unbiased estimate of $\btheta_*$ computed with $n$ samples in $\D$.

As $\wtheta_n$ is an unbiased estimate, we have that $\E_{\D}[Y_{n}]=\sum^A_{a=1}\bw(a)^{\top}\btheta_* = v(\pi)$.
Since we have an unbiased estimator of $v(\pi)$, minimizing the MSE is equivalent to minimizing the variance of $\min  \E_{\D}\big[\left(Y_n \!-\! \E[Y_n]\right)^2\big] = \min\E_{\D}\big[ \big(\sum_{a=1}^A\bw(a)^{\top}(\wtheta_n \!-\! \btheta_*)\big)^2\big]$

where the minimization is with respect to the data distribution $\D$, which is governed by observing rewards to actions determined by a behavior policy. In general, the behavior policy may be different from the target policy and it may even be non-stationary over the $n$ rounds of data collection.

\begin{align}
\bA_{\bb,\bSigma} &=\sum_{a\in\A}\bb(a)\big(\tfrac{\bx(a)}{\sigma(a)}\big)\big(\tfrac{\bx(a)}{\sigma(a)}\big)^\top %\nonumber\\
= \sum_{a\in\A}\bb(a)\tx(a)\tx(a)^\top \label{eq:design-matrix}
\end{align}

where, $\tx(a) = \bx(a)/\sigma(a)$.
Observe that our design matrix in (<ref>) captures the information about the action features $\bx(a)$, and variance $\sigma^2(a)$ and weights them by the sampling proportion $\bb(a)$. Then in the following proposition, we exactly characterize the loss with respect to the design matrix $\bA_{\bb,\bSigma}$, target policy $\pi(a)$ and action features $\bx(a)$. Define $\bV = \sum_a \bw(a)\bw(a)^\top$.

Let $\wtheta_n$ be the estimate of $\btheta_*$ after observing $n$ samples and define $\bw(a) = \pi(a)\bx(a)$. Define the design matrix as $\bA_{\bb,\bSigma} = \sum_{a=1}^A\bb(a)\left(\frac{\bx(a)}{\sigma(a)}\right)\left(\frac{\bx(a)}{\sigma(a)}\right)^\top$. Then the loss is given by

\begin{align*}
\E_\D\!\big[\big(\sum_{a=1}^A\bw(a)^\top(\wtheta_n-\btheta_{\star})\!\big)^{2}\big] \!\!=\!\! \frac{1}{n}\! \big(\!\sum_{a,a'}\bw(a)^\top\bA_{\bb,\bSigma}^{-1}\bw(a')\!\big) =:\L_n(\pi,\bb,\bSigma) .
\end{align*}

Proof (Overview)

The key idea is to show that the linear model yields for each action
$a \in[A], \widetilde{Y}_{n}(a) = \widetilde{\bx}_{n}(a)^{\top} \btheta^{\star}+\widetilde{\eta}_{n}(a)$ where we define

\begin{align*}
&\widetilde{Y}_{n}(a)=\sum_{t=1}^{T_{n}(a)} \tfrac{R_{t}^{}(a)}{\sigma(a) \sqrt{T_{n}(a)}}, \tx_n(a)=\tfrac{\sqrt{T_n(a)} \bx(a)}{ \sigma(a)},%\\
\widetilde{\eta}_{n}(a)=\sum_{t=1}^{T_n(a)} \tfrac{\eta_{t}^{}(a)}{\sigma(a) \sqrt{T_n(a)}}
\end{align*}

with $\eta_t(a)$ being the noise and $T_n(a)$ is the number of samples of action $a$. Next observe that using the independent noise assumption we have that  $\E[\widetilde{\eta}_{n}(a)]=0$ and $\Var\left[\widetilde{\eta}_{n}(a)\right]=1$. Let $\bX=$ $\left(\widetilde{\bx}_{n}(1)^{\top}, \cdots, \widetilde{\bx}_{n}(A)^{\top}\right)^{\top} \in \mathbb{R}^{A \times d}$ the induced design matrix of the policy and $\mathbf{Y} = [\widetilde{Y}_{n}(1), \widetilde{Y}_{n}(2),\ldots, \widetilde{Y}_{n}(A)]^\top$. The above ordinary least squares (OLS) problem has an optimal unbiased estimator $\wtheta_n=\left(\bX^{\top} \bX\right)^{-1} \bX^{\top} \mathbf{Y}$ [Fontaine et al., 2021]. Substituting the definition of $\wtheta_n$ yields the desired expression of the loss as stated in the proposition. The detailed proof is given in <Ref>. $\blacksquare$

Observe that the loss in our setting depends on the inverse of the design matrix denoted by $\bA^{-1}_{\bb,\bSigma}$, the target policy, as well as action features of the cross product of all arms $a,a'\in\A$. Hence, minimizing the loss is equivalent to minimizing the quantity $1/n (\sum_{a,a'}\bw(a)^\top\bA_{\bb,\bSigma}^{-1}\bw(a'))$. We call this the design. Further note that this design is different than a number of different prior notions of optimality such as D-, E-, T-, or G-optimality [Pukelsheim, 2006, Fedorov, 2013, Jamieson and Jain, 2022]. None of these previously proposed designs capture the objective of minimal MSE policy evaluation.

For example, G-optimality (as used by [Katz-Samuels et al., 2020, Mason et al., 2021, Katz-Samuels et al., 2021]) minimizes the worst-case error of $\max_{\bx(a)}\E_{\D}[(\bx(a)^\top(\wtheta_n-\btheta_*))^2]$ by minimizing the quantity $\max_{\bx(a)}\bx(a)^{\top}\bA^{-1}_{\bb}\bx(a)$ for homoscedastic noise. The E-optimal design minimizes $\max_{\|\mathbf{u}\|\leq 1}\E_{\D}[(\mathbf{u}^\top(\wtheta_n-\btheta_*))^2]$ by minimizing the minimum eigenvalue of the inverse of design matrix  [Mukherjee et al., 2022] and the A-optimal design minimizes $\E_{\D}[(\wtheta_n - \btheta_*)^2]$ by minimizing the trace of the inverse of design matrix [Fontaine et al., 2021].

We now state a few more notations for ease of exposition. Using <Ref> we define the  optimal behavior policy when the co-variance matrix $\bSigma_*$ is known as

\begin{align}
\bb^* \coloneqq \argmin_{\bb}\L_n(\pi, \bb, \bSigma_*). \label{eq:opt-oracle-sol}
\end{align}

where the loss $\L_n(\pi, \bb, \bSigma_*)$ is defined in <Ref>.
We define the optimal loss (with the knowledge of $\bSigma_*$)  as:

\begin{align}
\L_n(\pi, \bb^*, \bSigma_*) = \min_{\bb}\L_n(\pi, \bb, \bSigma_*). \label{eq:opt-oracle-loss}
\end{align}

Now an agnostic algorithm does not know the true co-variance matrix $\bSigma_*$ and must estimate the covariance matrix $\wSigma_\Gamma$ after conducting exploration for $\Gamma$ rounds.

Define the optimal

behavior policy for an arbitrary co-variance matrix $\wSigma_\Gamma$ and target policy $\pi$ as $\wb^* = \argmin_{\bb}\L_n(\pi, \bb, \wSigma_\Gamma)$

and the agnostic loss as

\begin{align}
\L_n(\pi, \wb^*, \wSigma_\Gamma) = \min_{\bb}\L_n(\pi, \bb, \wSigma_\Gamma) \label{eq:opt-agnostic-loss}
\end{align}

§.§ Computation of $\L_n(\pi, \bb, \bSigma)$

In this section, we digress a bit to discuss the computational aspect of $\L_n(\pi, \bb, \bSigma)$. Note that our loss function leads to a new type of design (design) and so the natural question to ask is how to optimize this loss function w.r.t. $\bb$? We show

in <Ref> that the loss $\L_n(\pi, \bb, \bSigma)$ for any arbitrary design proportion $\bb\in\triangle(\A)$  is strictly convex with respect to the proportion $\bb$. The proposition and its proof are given in <Ref>. Next in <Ref> we show that the gradient of the loss function is bounded. The proposition and its proof are given in <Ref>. We first state an assumption  that $\lambda_{\min }\left(\sum_{a=1}^A \bw(a) \bw(a)^{\top}\right)>0$ which is required for proving <Ref>.

(Distribution of $\pi$)
We assume that the set of actions $a$ such that $\!\pi(a) \!\!>\!\! 0$, spans $\R^d$ and $\R^{d\times d}$.

Note that this is a realistic assumption as if the target policy never takes an action that is needed to cover some dimension then we do not need to identify $\btheta_*$ in that dimension.

Using <Ref> and <Ref> we can effectively solve the design with gradient descent approaches [Lacoste-Julien and Jaggi, 2013, Berthet and Perchet, 2017].
We capture this convergence guarantee with the assumption of the existence of an approximation oracle.

(Approximation Oracle)

We assume access to an approximation oracle. Given a convex loss function $\L_n(\pi,\bb ,\bSigma)$ with minimizer $\bb^*$, the approximation oracle returns a proportion
$\wb^* = \argmin_{\bb} \L_n(\pi,\bb ,\bSigma)$ such that

$|\L_n(\pi,\wb^* ,\bSigma) - \L_n(\pi,\bb^* ,\bSigma)| \leq \epsilon$.

Therefore from <Ref>, and <ref> and using <Ref>, and <ref> we can get a computationally efficient solution to $\L_n(\pi, \bb, \bSigma)$. In the next section, we discuss the loss of an oracle.

§.§ Oracle Loss

Recall from <Ref>, that our final goal is to control the regret that compares the loss of an agnostic algorithm against the loss of an oracle.

In this section, we develop our theory for optimal data collection by considering an oracle for the heteroscedastic linear bandit setting.

We consider an oracle that has knowledge of $\bSigma_*$ but does not know $\btheta_*$.

Our goal is to identify the sampling proportion $\bb^*$ that such an oracle would select such that taking actions according to $\bb^*$ minimizes the loss.

After observing $n$ samples, let the weighted least square estimate be:

\begin{align}
\wtheta_{n} \coloneqq \arg\min_{\btheta}\sum_{t=1}^{n}\frac{1}{\sigma^{2}(I_{t})}(R(I_t)-\bv(I_t)^{\top}\btheta)^{2} \label{eq:weighted-least-square}
\end{align}

where $I_t$ is the action sampled at round $t$ and $\sigma^2(I_t)$ is the variance of  action $I_t$. Also note that this is an unbiased estimator of $\btheta_*$ (see <Ref>). We prove in <Ref> (see <Ref>) that if the oracle estimates $\wtheta_n$ using the weighted least square estimate in  (<ref>) then it can minimize the loss $\L^*_n(\pi, \bb^*, \bSigma_*)$.

Now the oracle uses the weighted least square estimate in (<ref>) to estimate $\wtheta_n$, knows $\bSigma_*$, and has access to the oracle approximator to solve the design. Then the following proposition bounds the loss of the oracle after $n$ samples.

(Oracle Loss)
Let the oracle sample each action $a$ for $\lceil n \bb^*(a)\rceil$ times, where $\bb^*$ is the solution to (<ref>). Define $\lambda_1(\bV)$ as the maximum eigenvalue of $\sum_{a,a'}\bw(a)\bw(a')^{\top}$. Then the loss is given by $
\L^*_n(\pi, \bb^*, \bSigma_*) \leq O\left(\tfrac{ d\lambda_1(\bV)
\log n}{n}\right) + O\left(\tfrac{1}{n}\right).

Proof (Overview) Note that the oracle knows the $\bSigma_*$ and uses $\wtheta_n$ in (<ref>) to estimate $\btheta_*$. We use <Ref> to show that $\L_n(\pi, \bb^*,\bSigma_*) \leq \lambda_1(\bV) d$ where $\bV = \sum_{a,a'}\bw(a)\bw(a')^\top$.

The proof follows by showing that $(\sum_{a=1}^A \bw(a)^\top (\wtheta_n-\btheta_{*}))^2$ is a sub-exponential variable.
Then using sub-exponential concentration inequality in <Ref> (<Ref>) and setting $\delta=O(1/n^2)$ we can bound the expected loss with high probability. The full proof is given in <Ref>. $\blacksquare$

Connection to prior work: We can interpret the result of <Ref> in the following way:

Consider the basic stochastic bandit setting which is a special case of our setting where $\bv(a)$ is a one-hot vector in $\mathbb{R}^A$.

In this case, from <Ref> we know that that $\bb^* = \arg\min_\bb \sum_a\frac{\pi^2(a)\sigma^2(a)}{\lceil \bb(a)n\rceil}$.

This captures the optimal number of times the actions should be pulled weighted by the target policy and their variance.

Solving for $\bb^*$, we obtain $\bb^*(a)\propto \pi^2(a)\sigma^2(a)$.

Note that this solution matches the optimal sampling proportion given by Antos et al., 2008, Carpentier and Munos, 2011, Carpentier and Munos, 2012, Carpentier et al., 2015 for the simple bandit setting.

Moreover, the loss of their oracle  decreases at the rate of $\widetilde{O}(An^{-1})$ whereas we decrease at the rate of $\widetilde{O}(d n^{-1})$. Our oracle loss scales with $d$ instead of $d^2$ as it knows the $\bSigma_*$ and does not need to explore $d^2$ directions.

So we obtain an equivalence between the solution of design and the solution from prior work in the basic bandit setting while considering a strictly more general setting.

§ AGNOSTIC ALGORITHM  AND REGRET ANALYSIS

In this section, we first present the agnostic algorithm and then analyze its regret.

§.§ Details of Algorithm

In practice, $\bSigma_*$ is unknown and so the oracle behavior policy cannot be directly computed.

Instead, we must first conduct a small amount of exploration to estimate $\bSigma_*$ and then use the estimate in place of $\bSigma_*$ in (<ref>).

Specifically, we define the forced exploration phase as the first $\Gamma$ rounds in which the algorithm conducts exploration to estimate $\bSigma_*$.

To conduct forced exploration we apply Principal Component Analysis (PCA) on the feature matrix $\bX$ and choose the most significant $d$ directions (directions having the highest variance). Then we choose one random action from each of these $d$ significant directions and sample them uniform randomly for $\Gamma$ rounds.

Since the algorithm explores first and then uses the estimate to compute the design, it can be viewed as an explore-then-commit algorithm [Rusmevichientong and Tsitsiklis, 2010, Lattimore and Szepesvári, 2020].

As we consider a structured setting we call this algorithm Structured Policy Evaluation Experimental Design ().
After $\Gamma = \sqrt{n}$ rounds,  estimates the covariance matrix $\wSigma_\Gamma$ as follows:

\begin{align}
\hspace*{-1.4em}\wSigma_\Gamma \!=\! \!\min_{\mathbf{S}}\sum_{t=1}^{\Gamma}\!\! \big[\langle \bv(I_{t}) \bv(I_{t})^{\top}, \mathbf{S}\rangle\!-\!(R(I_{t})\!-\!\bv(I_{t})^{\top} \wtheta_{t})^{2}\big]^{2} \label{eq:wSigma-Gamma}
\end{align}

where $\wtheta_t$ is the ordinary least square estimate of $\btheta_*$.

A similar Similar covariance estimation technique has been proposed before for the active regression setting though only for the case when $\bSigma_*$ has rank 1 [Chaudhuri et al., 2017].

The estimate of the covariance matrix $\wSigma_\Gamma$ is then fed to the oracle optimizer (<Ref>) to compute the sampling proportion $\wb$.

Finally, actions are chosen according to $\wb$ for the remaining $n-\Gamma$ rounds and then $\wtheta_n$ is computed.

Full pseudocode is given in <Ref>.

Structured Policy Evaluation Experimental Design ()

[1]
Input: Finite set of actions $\A$, target policy $\pi$, budget $n$.
Conduct forced exploration for $\Gamma=\sqrt{n}$ rounds and estimate $\wSigma_\Gamma$ using (<ref>).

Let $\wb_{} \in \triangle(\A)$ be the

minimizer of $\L_n(\pi,\bb ,\wSigma_\Gamma)$.

Pull each action $a$ exactly $T_n(a) = \left\lfloor\wb_{}(a) (n-\Gamma) \right\rfloor$ times. Set  $\H(a) \coloneqq \{I_t,R(I_t)\}_{t=\Gamma}^{T_n(a)}$ by selecting $I_t=a$ according to $\wb(a)$. Set $\D\leftarrow \cup_a\H(a)$.
Construct the policy weighted least squares estimator $\wtheta_{n}$ using $\D$ using only the observations from step 4.
Output: $\D$ and $\wtheta_n$.

§.§ Regret Analysis of

In this section we first state our regret definition and then analyze the regret of the agnostic algorithm .

The regret for the agnostic algorithm for the estimated behavior policy $\wb$ is given by

\begin{align}
\cR_n &= \bL_n(\pi,\wb,\wSigma_\Gamma) - \L_n(\pi,\bb^*,\bSigma_*). \label{eq:regret-definition}
\end{align}

where $\bL_n(\pi,\wb,\wSigma_\Gamma)$ is the loss of the agnostic algorithm and

$\L_n(\pi,\bb^*,\bSigma_*)$ is the oracle loss defined in (<ref>).

We now state the main theorem for the regret of .

(Regret of <Ref>, informal)

The regret of <Ref> for $n\geq O(\tfrac{d^4\log^2 (A/\delta)}{\sigma^4_{\min}})$

running design in <ref> is given by

$\cR_n = O\left(\frac{d^3 \log(n )}{n^{3/2}}\right)$.

Discussion (Regret): <Ref> states that the regret of <Ref> scales as $O(d^3\log(n)/n^{3/2})$ where $d$ is the dimension of $\btheta^*$. Note that our regret bound depends on the underlying feature dimension $d$ instead of actions $A$, and scales as $\widetilde{O}(d^3 n^{-3/2})$ which gives a positive answer to the main question of whether such a result is possible.

In the case where $d^3 < A$, we have a tighter bound than Carpentier and Munos, 2011.
Furthermore, the result of Carpentier and Munos, 2011, Carpentier and Munos, 2012, Carpentier et al., 2015 cannot be easily extended to take advantage of structure in the linear bandit setting.

Finally, note that we improve upon the

A-optimal design studied in Fontaine et al., 2021, as their regret depends on actions $A$ and scales as $O(\frac{A\log n}{n^{3/2}})$.

Proof (Overview) of <Ref>: We now outline the key steps for proving <Ref>.

Step 1(Concentration of (OLS-WLS): We now state a concentration lemma that is key to proving the regret of .
This lemma is novel for our proof because we must estimate the underlying covariance matrix $\bSigma_*$ using OLS estimator for $\Gamma$ rounds. Then use the estimation $\wSigma_\Gamma$ in WLS estimator.

We define the variance concentration good event till $\Gamma$ using our forced exploration as:

\begin{align}
\hspace*{-1.6em}\xi^{var}_\delta(\Gamma) \!\coloneqq\!\! \big\{\forall a,|\bx(a)^{\top}(\wSigma_\Gamma - \bSigma_*)\bx(a)| < \tfrac{2C d^2 \log ({A}/{\delta})}{\Gamma}\big\} \label{eq:event-xi-delta}
\end{align}

(OLS-WLS Concentration Lemma)

After $\Gamma$ samples of exploration, we can show that $\Pb\left(\xi^{var}_\delta(\Gamma)\right)\geq  1 -8\delta$
where, $C > 0$ is a constant.

Proof (Overview) of <Ref>: We have an initial estimate $\wtheta_\Gamma$ of $\btheta^*$ and the squared residual $y_{t}:=(\bx_{t}^{\top}\wtheta_\Gamma-r_{t})^{2}$ that  obtains by estimating of $\bSigma_*$ via $
\min_{\bS\in\R^{d\times d}}\sum_{t=1}^{\Gamma}(\left\langle \bx_{t}\bx_{t}^{\top},\bS\right\rangle -y_{t})^{2}$. Let $\zeta_\Gamma \coloneqq \wtheta_\Gamma-\btheta^*$ then we can show that $y_{t}=\bx_{t}^{\top}\bSigma_* \bx_{t}+\epsilon_{t}$ and the noise $\epsilon_t$ can be bounded by

$\epsilon_t\leq \underbrace{2(\eta_{t}^{2}-\E\left[\eta_{t}^{2}\right])}_{\textbf{Part A}} + \underbrace{2\left(\bx_{t}^{\top}\zeta_\Gamma\right)^{2}}_{\textbf{Part B}}$.

For the part A, observe that $\eta^2_t$ is a sub-exponential random variable as  $\eta_t\sim\SG(0,\bx_{t}^{\top}\bSigma_* \bx_{t})$. Hence we can use sub-exponential concentration inequality from <Ref> (<Ref>) to bound it. For part B first recall that $\zeta_\Gamma \coloneqq \wtheta_\Gamma-\btheta^*$ and we use  <Ref> (<Ref>) to bound it. Combining the two parts give the desired concentration inequality. The proof is in <Ref>. $\blacksquare$

Step 2 (Agnostic loss $\bL_n(\pi,\wb,\wSigma_\Gamma)$): In this step we bound the agnostic loss $\bL_n(\pi,\wb,\wSigma_\Gamma)$.

The <Ref> leads to  <Ref> (<Ref>) which shows that for $n\geq 16C^2 d^4\log^2 (A/\delta)/ \sigma^4_{\min}$ we have that $\L_n(\pi, \wb, \wSigma_\Gamma) \leq (1+2C_\Gamma(\delta)) \sum_{a,a'}\bw(a)\bA_{\wb^*,\wSigma_\Gamma}^{-1}\bw(a')$, where $C_\Gamma(\delta) = \tfrac{2C d^2  \log (A / \delta)}{\sigma^2_{\min}\Gamma}$.

Note that Fontaine et al., 2021 does not require this approach as the variances of each action do not share a common structure. Similarly, this approach differs from the time-dependent variance model of Zhang et al., 2021, Zhao et al., 2022.

Now observe that for an agnostic algorithm, the loss function is defined as $\L_n(\pi, \wb, \wSigma_\Gamma)$. However,  $\L_n(\pi, \wb, \wSigma_\Gamma) \neq \sum_{a,a'}\bw(a)\bA_{\wb^*,\wSigma_\Gamma}^{-1}\bw(a')$.  We denote $\bL_n(\pi,\wb,\wSigma_\Gamma) \coloneqq \E\big[\big(\bz^{\top}(\wtheta_{n-\Gamma} - \btheta_*)\big)^2\big]$ where $\wtheta_{n-\Gamma}$ is the estimation of $\btheta_*$ after $n-\Gamma$ samples.

We now bound the quantity $\bL_n(\pi, \wb, \wSigma_\Gamma)$ using <Ref>. In it, we show that

\begin{align*}
\E\big[\big(\bz^{\top}(\wtheta_{n-\Gamma} - \btheta_*)\big)^2\leq \left(1+2C_\Gamma(\delta)\right)\bz^\top (\bX_{n-\Gamma}^{\top}\wSigma_{\Gamma}^{-1}\bX_{n-\Gamma})^{-1}\bz = \tfrac{\left(1+2C_\Gamma(\delta)\right)}{n - \Gamma}\sum_{a,a'}\bw(a)^\top\bA_{\wb^*, \wSigma_\Gamma}^{-1}\bw(a')
\end{align*}

Step 3 (Regret Decomposition):
Define the regret as in (<ref>). Then recall that $\bb^*\in\triangle(\A)$ is the optimal design in (<ref>) and $\wb^* \!\!=\! \argmin_{\bb}\L_n(\pi, \bb, \wSigma_\Gamma)$ is the agnostic design.

Define $\L'_{n-\Gamma}(\pi,\wb,\wSigma_\Gamma) \coloneqq \tfrac{\left(1+2C_\Gamma(\delta)\right)}{n - \Gamma}\sum_{a,a'}\bw(a)^\top\bA_{\wb^*, \wSigma_\Gamma}^{-1}\bw(a')$. Then we can show that the regret can be decomposed into three parts:

\begin{align*}
\cR_n &= \underbrace{\L'_{n-\Gamma}(\pi,\wb,\wSigma_\Gamma) - \L'_{n-\Gamma}(\pi,\wb^*,\wSigma_\Gamma)}_{\textbf{Approximation error}}+ \underbrace{\L'_{n-\Gamma}(\pi,\wb^*,\wSigma_\Gamma) -  \L_n(\pi,\bb^*,\wSigma_\Gamma)}_{\textbf{Comparing two diff loss}}\\
&\quad + \underbrace{\L_n(\pi,\bb^*,\wSigma_\Gamma) - \L_n(\pi,\bb^*,\bSigma_*)}_{\textbf{Estimation error of $\bSigma_*$}}.
\end{align*}

For the approximation error we need access to an oracle (<Ref>) that gives $\epsilon$ approximation error. Then setting $\epsilon=\tfrac{1}{\sqrt{n}}$ we have that the estimation error is upper bounded by $n^{-3/2}$.

For comparing two different loss parts, we use the definition of $C_\Gamma(\delta)$ to bound it as $O(\tfrac{d^2\log(A/\delta)}{n^{3/2}})$ as shown in (<ref>) in <Ref>.

Now observe that the third quantity (estimation error of $\bSigma_*$) contains $\L_n(\pi, \bb^*, \wSigma_\Gamma)$ that depends on the design matrix $\bA^{-1}_{\bb^*,\wSigma_\Gamma}$ which in turn depends on the estimation of $\wSigma_\Gamma$.

Similarly $\L_n(\pi, \bb^*, \bSigma_*)$ in the third quantity depends on the design matrix $\bA^{-1}_{\bb^*,\bSigma_*}$ which in turn depends on the true $\bSigma_*$.

Hence, we now bound the concentration of the loss under $\bA^{-1}_{\bb^*,\wSigma_\Gamma}$ against the design matrix $\bA^{-1}_{\bb^*,\bSigma_*}$ in the following lemma.

(Loss Concentration of design matrix) Let $\wSigma_\Gamma$ be the empirical estimate of $\bSigma_*$, and  $\bV=\sum_{a,a'}\bw(a)\bw(a')^{\top}$. We have that for any arbitrary proportion $\bb$ the following

\begin{align*}
\Pb&\big(\big|\sum_{a,a'} \bw(a)^{\top}(\bA^{-1}_{\bb^*, \wSigma_\Gamma} - \bA^{-1}_{\bb^*, \bSigma_*})\bw(a')\big|%\\
\leq \tfrac{2C B^* d^3 \log (A/ \delta)}{\Gamma}\big)\geq 1-\delta
\end{align*}

where $B^*$ is a problem-dependent quantity

and $C>0$ is a universal constant.

Proof (Overview) of <Ref>: We can decompose

$|\sum_{a,a'} \!\bw(a)^{\top}(\bA^{-1}_{\bb^*, \wSigma_\Gamma}\!-\!\bA^{-1}_{\bb^*, \bSigma_*})\bw(a')| %\\
\!\leq\! \|\bu\|\!\underbrace{\left\|\bA_{\bb^*,\bSigma_*} - \bA_{\bb^*,\wSigma_\Gamma}\right\|}_{\Delta}\!\|\mathbf{v}\| $

where, $\|\bu\|=\|\bA^{-1}_{\bb^*, \bSigma_*} \bw\|$ and $\|\mathbf{v}\| = \|\bA^{-1}_{\bb^*, \wSigma_\Gamma}\bw\|$. First, observe that $\|\bu\|$ is a problem-dependent quantity. Then to bound $\Delta$ we bound the $\|\wSigma_\Gamma - \bSigma_*\| \leq \frac{2C d^2 \log (A / \delta)}{\Gamma}$. Finally to bound $\|\mathbf{v}\|$ we need to bound $\wsigma^2_\Gamma(a) \leq \sigma^2(a) + \frac{2C d^2 \log (A / \delta)}{\Gamma}$ where $\wsigma^2_\Gamma(a)$ is the empirical variance of $\sigma^2(a)$. Combining everything yields the desired result. The proof is in <Ref> $\blacksquare$

One of our key technical contributions in <Ref> is to show that the concentration of the two losses $\L_n(\pi, \bb^*, \wSigma_\Gamma)$, and $\L_n(\pi, \bb^*, \bSigma_*)$ scales with $d^3$ instead of the number of actions $A$. In contrast a similar loss concentration in Fontaine et al., 2021 scales with $A$.

Now using <Ref>, setting the exploration factor $\Gamma=\sqrt{n}$, and $\delta=\frac{1}{n}$ we can show that the estimation error is upper bounded by

$\frac{B^*Cd^3\log (n )}{n^{3/2}}  + \frac{d^2}{n^2}\Tr(\sum_{a,a'}\bw(a)\bw(a')^\top)$.

Combining everything we have the regret of  as $ O(\tfrac{B^*d^3\log (n )}{n^{3/2}})$. The full proof of <Ref> is in <Ref>. $\blacksquare$.

<Ref> upper bounds the regret of our agnostic algorithm  compared to an oracle algorithm with knowledge of $\bSigma_*$.

To quantify the tightness of our upper bound, we now turn to whether we can lower bound the regret of .

For our final theoretical result, we consider a slightly different notion of regret:
\cR'_n \coloneqq \L_n(\pi, \wb, \bSigma_*) - \L_n(\pi, \bb^*, \Sigma_*).

This notion of regret captures how sub-optimal the estimated $\wb$ is compared to $\bb^*$ and not additional error incurred by using an estimate of $\Sigma^*$ in the PWLS estimator.

We conjecture that $\cR'_n$ is indeed a lower bound to $\cR_n$ as we have established in <Ref> that the minimum variance estimator is the PWLS estimator using $\Sigma^\star$.

Intuitively, $\L_{n}(\pi,\wb,\bSigma_*)$ is a lower bound to $\overline{\L}_n(\pi, \wb, \widehat{\Sigma}_\Gamma)$ as estimation error will likely only increase when using $\widehat{\Sigma}_\Gamma$ in place of $\Sigma_*$ in the PWLS estimator.

We leave proving that $\cR'_n$ is a lower bound to $\cR_n$ in future work.

2(Lower Bound)

Let $|\bTheta| \!\!=\!\! 2^d$, $\btheta^*\!\in\!\bTheta$. Then any $\delta$-PAC policy $\bb$
satisfies $\cR'_{n} \!\!=\!\! \L_{n}(\pi,\wb,\bSigma_*) \!-\! \L_{n}(\pi,\bb^*,\bSigma_*) \!\geq\! \Omega\left(\frac{d^2\lambda_d(\bV)\log({n})}{{n}^{3/2}}\right)$ for the environment
in (<ref>).

Proof (Overview:) The proof follows the change of measure argument [Lattimore and Szepesvári, 2020]. We follow the proof technique of Huang et al., 2017, Mukherjee et al., 2022. We reduce our linear bandit problem to the hypothesis testing setting and state a worst-case environment as in (<ref>). We then show that the regret of any $\delta$-PAC algorithm against an oracle in this environment must scale as $\Omega(\log n/n^{3/2})$. The proof is given in <Ref>. $\blacksquare$

From the above result, the upper bound of  regret $\cR_n$ matches the lower bound of regret $\cR'_n$ in $n$ but suffers an additional factor of $d$.

§ EXPERIMENTS

We now conduct numerical experiments to show that  decreases MSE faster than other baselines.

As baselines, we compare against , , [Fontaine et al., 2021], [Wan et al., 2022].

The  algorithm simply runs the target policy to collect data and evaluate $\pi$, whereas the (as discussed in <Ref>) is similar to  but knows $\bSigma_*$.

Of these baselines, , and are the closest in relation to our work.

We experiment with design because this criterion minimizes the average variance of the estimates of the regression coefficients and is most closely aligned with our goal.

Note that Fontaine et al., 2021 considers policy improvement. The work of Wan et al., 2022 considers data collection under safety constraints using Inverse Propoensity Weighting. The  implement this variant. In our unconstrained policy evaluation setting their approach boils down to just G-optimal design [Pukelsheim, 2006].

Further experimental details are in <Ref>.

(Left) MSE plot for the Unit ball. (Middle-left) MSE plot for the Movielens dataset. (Middle-Right) MSE plot for Red Wine Quality dataset. (Right) MSE plot for Air Quality dataset. The vertical axis gives MSE and the horizontal axis is the number of rounds.
The vertical axis is log-scaled and confidence bars show one standard error.

Unit Ball: We perform this experiment on a set of $5$ actions that are arranged in a unit ball in $\R^2$ to show that  allocates proportion to the most informative action (weighted by their variance).

<Ref> (Left) shows that   outperforms , , and .

We also show  in this setting to show how quickly  converges to it. However, for settings based on real-life data, we do not have such oracles.

Movielens Dataset:

Consider a startup that wants to recommend movies to users based on their ratings.

They have access to a target policy and want to evaluate it on a limited informative dataset before deploying it for full public use.

We use real-world Movielens 1M dataset [Lam and Herlocker, 2016] datasets for this experiment.

We apply low-rank factorization to the rating matrix to obtain $5$-dimensional representations of users and movies.
We then fit a weighted least square estimate of $\btheta^*$ and $\bSigma_*$.

We generate the reward using this $\btheta^*$ and $\bSigma_*$.

Then we use  and other baselines to generate the small informative dataset to evaluate the target policy and this experiment is shown in <Ref> (Left).

initially conducts forced exploration to estimate $\btheta_*$, $\bSigma_*$ and incurs slightly higher MSE but the MSE decreases faster than other baselines as the number of rounds increases.

Red Wine Quality: We use a similar motivation as before to conduct this experiment. However, we now consider an online wine company that wants to recommend wines to users and wants to evaluate a target policy before full deployment.

We perform this  experiment on real-world dataset Red Wine Quality from UCI datasets [Cortez et al., 2009].

The dataset consists of $1600$ samples (actions) of red wine with each sample $a$ having feature $\bx(a)\in\mathbb{R}^{11}$ and their ratings. We fit a weighted least square estimate to the original dataset and get an estimate of $\btheta^*$ and $\bSigma_*$.

We generate the reward using this $\btheta^*$ and $\bSigma_*$.

Then we use  to generate the informative dataset to evaluate the target policy.

<Ref> (Middle-Right) shows that   outperforms other baselines as horizon increases.

Air Quality: We now consider a setting where a government agency wants to record air quality and notify the public. However, it wants to evaluate a target policy on a limited informative dataset before full deployment.

We perform this experiment on real-world dataset Air-Quality from UCI datasets [De Vito et al., 2008].

The dataset consists of $1500$ samples (actions) with each sample $a$ having feature $\bx(a)\in\mathbb{R}^{6}$ and their air quality value. We fit a weighted least square estimate to the original dataset and get an estimate of $\btheta^*$ and $\bSigma_*$. We generate the reward using this $\btheta^*$ and $\bSigma_*$.

Then we use  and other baselines (which do not know $\btheta^*$ and $\bSigma_*$) to generate the informative dataset to evaluate the target policy and this experiment is shown in <Ref> (Right)

<Ref> (Bottom-Right) shows that  MSE decreases faster than other baselines as the number of rounds increases.

§ CONCLUSIONS AND FUTURE DIRECTIONS

In this paper, we proposed  for optimal data collection for policy evaluation in linear bandits with  heteroscedastic reward noise. We formulated a novel optimal design problem, design, for which the optimal behavior policy is the solution that will produce minimal MSE policy evaluation when using a weighted least square estimate of the hidden reward parameters $\btheta_*$ and $\bSigma_*$.

We showed the regret of  degrades at the rate of $\widetilde{O}(d^3n^{-3/2})$ and matches the lower bound of $\widetilde{O}(d^2n^{-3/2})$ except a factor of $d$.

In contrast the  suffers a regret of $\widetilde{O}(n^{-1})$ [Carpentier et al., 2015].

We showed empirically that our design outperforms other optimal designs.

In future work, we intend to extend the result to a more general class of hard problems such as collecting data to minimize the MSE of multiple target policies under a generalized linear bandit setting.

[Abbasi-Yadkori et al., 2011]
Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári.
Improved algorithms for linear stochastic bandits.
Advances in neural information processing systems, 24, 2011.

[Agrawal and Goyal, 2012]
Shipra Agrawal and Navin Goyal.
Analysis of thompson sampling for the multi-armed bandit problem.
In Conference on learning theory, pages 39–1. JMLR Workshop
and Conference Proceedings, 2012.

[Antos et al., 2008]
András Antos, Varun Grover, and Csaba Szepesvári.
Active learning in multi-armed bandits.
In International Conference on Algorithmic Learning Theory,
pages 287–302. Springer, 2008.

[Auer et al., 2002]
Peter Auer, Nicolò Cesa-Bianchi, and Paul Fischer.
Finite-time Analysis of the Multiarmed Bandit Problem.
Machine Learning, 470 (2):0 235–256, May
ISSN 1573-0565.
URL <https://doi.org/10.1023/A:1013689704352>.

[Berthet and Perchet, 2017]
Quentin Berthet and Vianney Perchet.
Fast rates for bandit optimization with upper-confidence frank-wolfe.
Advances in Neural Information Processing Systems, 30, 2017.

[Bottou et al., 2013]
Léon Bottou, Jonas Peters, Joaquin Quiñonero-Candela, Denis X Charles,
D Max Chickering, Elon Portugaly, Dipankar Ray, Patrice Simard, and
Ed Snelson.
Counterfactual reasoning and learning systems: The example of
computational advertising.
Journal of Machine Learning Research, 140 (11), 2013.

[Bouchard et al., 2016]
Guillaume Bouchard, Théo Trouillon, Julien Perez, and Adrien Gaidon.
Online learning to sample.
arXiv preprint arXiv:1506.09016, 2016.

[Cai et al., 2021]
Hengrui Cai, Chengchun Shi, Rui Song, and Wenbin Lu.
Deep jump learning for off-policy evaluation in continuous treatment
Advances in Neural Information Processing Systems,
34:0 15285–15300, 2021.

[Carpentier and Munos, 2011]
Alexandra Carpentier and Rémi Munos.
Finite-time analysis of stratified sampling for monte carlo.
In NIPS-Twenty-Fifth Annual Conference on Neural Information
Processing Systems, 2011.

[Carpentier and Munos, 2012]
Alexandra Carpentier and Rémi Munos.
Minimax number of strata for online stratified sampling given noisy
In International Conference on Algorithmic Learning Theory,
pages 229–244. Springer, 2012.

[Carpentier et al., 2015]
Alexandra Carpentier, Remi Munos, and András Antos.
Adaptive strategy for stratified monte carlo sampling.
J. Mach. Learn. Res., 16:0 2231–2271, 2015.

[Chaudhuri et al., 2017]
Kamalika Chaudhuri, Prateek Jain, and Nagarajan Natarajan.
Active heteroscedastic regression.
In International Conference on Machine Learning, pages
694–702. PMLR, 2017.

[Ciosek and Whiteson, 2017]
Kamil Ciosek and Shimon Whiteson.
OFFER: Off-environment reinforcement learning.
In Proceedings of the 31st AAAI Conference on Artificial
Intelligence (AAAI), 2017.

[Cortez et al., 2009]
Paulo Cortez, António Cerdeira, Fernando Almeida, Telmo Matos, and José
Modeling wine preferences by data mining from physicochemical
Decision support systems, 470 (4):0 547–553,

[De Vito et al., 2008]
Saverio De Vito, Ettore Massera, Marco Piga, Luca Martinotto, and Girolamo
Di Francia.
On field calibration of an electronic nose for benzene estimation in
an urban pollution monitoring scenario.
Sensors and Actuators B: Chemical, 1290 (2):0
750–757, 2008.

[Dudík et al., 2014]
Miroslav Dudík, Dumitru Erhan, John Langford, and Lihong Li.
Doubly robust policy evaluation and optimization.

[Fang et al., 1994]
Yuguang Fang, Kenneth A Loparo, and Xiangbo Feng.
Inequalities for the trace of matrix product.
IEEE Transactions on Automatic Control, 390
(12):0 2489–2490, 1994.

[Fedorov, 2013]
Valerii Vadimovich Fedorov.
Theory of optimal experiments.
Elsevier, 2013.

[Fontaine et al., 2021]
Xavier Fontaine, Pierre Perrault, Michal Valko, and Vianney Perchet.
Online a-optimal design and active linear regression.
In International Conference on Machine Learning, pages
3374–3383. PMLR, 2021.

[Greene, 2002]
William H Greene.
000. econometric analysis, 2002.

[Hanna et al., 2017]
Josiah P. Hanna, Philip S. Thomas, Peter Stone, and Scott Niekum.
Data-Efficient Policy Evaluation Through Behavior Policy
arXiv:1706.03469 [cs], June 2017.
URL <http://arxiv.org/abs/1706.03469>.
arXiv: 1706.03469.

[Huang et al., 2017]
Ruitong Huang, Mohammad M. Ajallooeian, Csaba Szepesvári, and Martin
Structured best arm identification with fixed confidence.
In Steve Hanneke and Lev Reyzin, editors, International
Conference on Algorithmic Learning Theory, ALT 2017, 15-17 October 2017,
Kyoto University, Kyoto, Japan, volume 76 of Proceedings of Machine
Learning Research, pages 593–616. PMLR, 2017.
URL <http://proceedings.mlr.press/v76/huang17a.html>.

[Jamieson and Jain, 2022]
Kevin Jamieson and Lalit Jain.
Interactive machine learning.

[Kallus et al., 2021]
Nathan Kallus, Yuta Saito, and Masatoshi Uehara.
Optimal off-policy evaluation from multiple logging policies.
In International Conference on Machine Learning, pages
5247–5256. PMLR, 2021.

[Katz-Samuels et al., 2020]
Julian Katz-Samuels, Lalit Jain, Kevin G Jamieson, et al.
An empirical process approach to the union bound: Practical
algorithms for combinatorial and linear bandits.
Advances in Neural Information Processing Systems,
33:0 10371–10382, 2020.

[Katz-Samuels et al., 2021]
Julian Katz-Samuels, Jifan Zhang, Lalit Jain, and Kevin Jamieson.
Improved algorithms for agnostic pool-based active classification.
In International Conference on Machine Learning, pages
5334–5344. PMLR, 2021.

[Kiefer and Wolfowitz, 1960]
Jack Kiefer and Jacob Wolfowitz.
The equivalence of two extremum problems.
Canadian Journal of Mathematics, 12:0 363–366, 1960.

[Kirschner and Krause, 2018]
Johannes Kirschner and Andreas Krause.
Information directed sampling and bandits with heteroscedastic noise.
In Conference On Learning Theory, pages 358–384. PMLR, 2018.

[Kohavi and Longbotham, 2017]
Ron Kohavi and Roger Longbotham.
Online controlled experiments and a/b testing.
Encyclopedia of machine learning and data mining, 70
(8):0 922–929, 2017.

[Lacoste-Julien and Jaggi, 2013]
Simon Lacoste-Julien and Martin Jaggi.
An affine invariant linear convergence analysis for frank-wolfe
arXiv preprint arXiv:1312.7864, 2013.

[Lai and Robbins, 1985]
T. L Lai and Herbert Robbins.
Asymptotically efficient adaptive allocation rules.
Advances in Applied Mathematics, 60 (1):0
4–22, March 1985.
ISSN 0196-8858.

[Lam and Herlocker, 2016]
Shyong Lam and Jon Herlocker.
MovieLens Dataset.
http://grouplens.org/datasets/movielens/, 2016.

[Lattimore and Szepesvári, 2020]
Tor Lattimore and Csaba Szepesvári.
Bandit algorithms.
Cambridge University Press, 2020a.

[Lattimore and Szepesvári, 2020]
Tor Lattimore and Csaba Szepesvári.
Bandit algorithms.
Cambridge University Press, 2020b.

[Li et al., 2011]
Lihong Li, Wei Chu, John Langford, and Xuanhui Wang.
Unbiased offline evaluation of contextual-bandit-based news article
recommendation algorithms.
In Proceedings of the fourth ACM international conference on
Web search and data mining, pages 297–306, 2011.

[Li et al., 2015]
Lihong Li, Rémi Munos, and Csaba Szepesvári.
Toward minimax off-policy value estimation.
In Artificial Intelligence and Statistics, pages 608–616.
PMLR, 2015.

[Mason et al., 2021]
Blake Mason, Romain Camilleri, Subhojyoti Mukherjee, Kevin Jamieson, Robert
Nowak, and Lalit Jain.
Nearly optimal algorithms for level set estimation.
arXiv preprint arXiv:2111.01768, 2021.

[Mukherjee et al., 2022]
Subhojyoti Mukherjee, Josiah P Hanna, and Robert Nowak.
Revar: Strengthening policy evaluation via reduced variance sampling.
arXiv preprint arXiv:2203.04510, 2022a.

[Mukherjee et al., 2022]
Subhojyoti Mukherjee, Ardhendu S Tripathy, and Robert Nowak.
Chernoff sampling for active testing and extension to active
In International Conference on Artificial Intelligence and
Statistics, pages 7384–7432. PMLR, 2022b.

[Oosterhuis and de Rijke, 2020]
Harrie Oosterhuis and Maarten de Rijke.
Taking the Counterfactual Online: Efficient and Unbiased
Online Evaluation for Ranking.
Proceedings of the 2020 ACM SIGIR on International Conference
on Theory of Information Retrieval, pages 137–144, September 2020.
URL <http://arxiv.org/abs/2007.12719>.
arXiv: 2007.12719.

[Pukelsheim, 2006]
Friedrich Pukelsheim.
Optimal design of experiments.
SIAM, 2006.

[Rigollet and Hütter, 2015]
Phillippe Rigollet and Jan-Christian Hütter.
High dimensional statistics.
Lecture notes for course 18S997, 8130 (814):0
46, 2015.

[Riquelme et al., 2017]
Carlos Riquelme, Mohammad Ghavamzadeh, and Alessandro Lazaric.
Active learning for accurate estimation of linear models.
In International Conference on Machine Learning, pages
2931–2939. PMLR, 2017.

[Rusmevichientong and Tsitsiklis, 2010]
Paat Rusmevichientong and John N Tsitsiklis.
Linearly parameterized bandits.
Mathematics of Operations Research, 350 (2):0
395–411, 2010.

[Su et al., 2020]
Yi Su, Maria Dimakopoulou, Akshay Krishnamurthy, and Miroslav Dudík.
Doubly robust off-policy evaluation with shrinkage.
In International Conference on Machine Learning, pages
9167–9176. PMLR, 2020.

[Sutton and Barto, 2018]
Richard S Sutton and Andrew G Barto.
Reinforcement learning: An introduction.
MIT press, 2018.

[Swaminathan et al., 2017]
Adith Swaminathan, Akshay Krishnamurthy, Alekh Agarwal, Miro Dudik, John
Langford, Damien Jose, and Imed Zitouni.
Off-policy evaluation for slate recommendation.
Advances in Neural Information Processing Systems, 30, 2017.

[Thompson, 1933]
William R. Thompson.
On the likelihood that one unknown probability exceeds another in
view of the evidence of two samples.
Biometrika, 250 (3-4):0 285–294, December
ISSN 0006-3444.
URL <https://doi.org/10.1093/biomet/25.3-4.285>.

[Tsybakov, 2008]
Alexandre B Tsybakov.
Introduction to nonparametric estimation.
Springer Science & Business Media, 2008.

[Tucker and Joachims, 2022]
Aaron David Tucker and Thorsten Joachims.
Variance-Optimal Augmentation Logging for Counterfactual
Evaluation in Contextual Bandits.
arXiv:2202.01721 [cs], February 2022.
URL <http://arxiv.org/abs/2202.01721>.
arXiv: 2202.01721.

[Wainwright, 2019]
Martin J Wainwright.
High-dimensional statistics: A non-asymptotic viewpoint,
volume 48.
Cambridge university press, 2019.

[Wan et al., 2022]
Runzhe Wan, Branislav Kveton, and Rui Song.
Safe exploration for efficient policy evaluation and comparison.
arXiv preprint arXiv:2202.13234, 2022.

[Wang et al., 2017]
Yu-Xiang Wang, Alekh Agarwal, and Miroslav Dudık.
Optimal and adaptive off-policy evaluation in contextual bandits.
In International Conference on Machine Learning, pages
3589–3597. PMLR, 2017.

[Whittle, 1958]
P Whittle.
A multivariate generalization of tchebichev's inequality.
The Quarterly Journal of Mathematics, 90 (1):0
232–240, 1958.

[Zhang et al., 2021]
Zihan Zhang, Jiaqi Yang, Xiangyang Ji, and Simon S Du.
Improved variance-aware confidence sets for linear bandits and linear
mixture mdp.
Advances in Neural Information Processing Systems,
34:0 4342–4355, 2021.

[Zhao et al., 2022]
Heyang Zhao, Dongruo Zhou, Jiafan He, and Quanquan Gu.
Bandit learning with general function classes: Heteroscedastic noise
and variance-dependent regret bounds.
arXiv preprint arXiv:2202.13603, 2022.

[Zhou and Gu, 2022]
Dongruo Zhou and Quanquan Gu.
Computationally efficient horizon-free reinforcement learning for
linear mixture mdps.
arXiv preprint arXiv:2205.11507, 2022.

[Zhou et al., 2021]
Dongruo Zhou, Quanquan Gu, and Csaba Szepesvari.
Nearly minimax optimal reinforcement learning for linear mixture
markov decision processes.
In Conference on Learning Theory, pages 4532–4576. PMLR,

[Zhou et al., 2017]
Xin Zhou, Nicole Mayer-Hamblett, Umer Khan, and Michael R Kosorok.
Residual weighted learning for estimating individualized treatment
Journal of the American Statistical Association, 1120
(517):0 169–187, 2017.

[Zhu and Kveton, 2021]
Ruihao Zhu and Branislav Kveton.
Safe data collection for offline and online policy learning.
arXiv preprint arXiv:2111.04835, 2021.

[Zhu and Kveton, 2022]
Ruihao Zhu and Branislav Kveton.
Safe optimal design with applications in off-policy learning.
In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera,
editors, Proceedings of The 25th International Conference on Artificial
Intelligence and Statistics, volume 151 of Proceedings of Machine
Learning Research, pages 2436–2447. PMLR, 28–30 Mar 2022.
URL <https://proceedings.mlr.press/v151/zhu22a.html>.

§ APPENDIX

§.§ Related Works and Motivations

§.§ Probability Tools

[Kiefer and Wolfowitz, 1960]
Assume that $\A \subset \mathbb{R}^{d}$ is compact and $\operatorname{span}(\A)=\mathbb{R}^{d}$. Let $\pi: \A \rightarrow[0,1]$ be a distribution on $\A$ so that $\sum_{a \in \A} \pi(a)=1$ and $\bV(\pi) \in \mathbb{R}^{d \times d}$ and $g(\pi) \in \mathbb{R}$ be given by

\begin{align*}
\bV(\pi)=\sum_{a \in \A} \pi(a) a a^{\top}, \quad g(\pi)=\max _{a \in \A}\|a\|_{\tX(\pi)^{-1}}^{2}
\end{align*}

Then the following are equivalent:

* $\pi^{*}$ is a minimizer of $g$.

* $\pi^{*}$ is a maximizer of $f(\pi)=\log \det \bV(\pi)$.

* $g\left(\pi^{*}\right)=d$.

Furthermore, there exists a minimizer $\pi^{*}$ of $g$ such that $\left|\operatorname{Supp}\left(\pi^{*}\right)\right| \leq d(d+1) / 2$.

(Sub-Exponential Concentration)
Suppose that $X$ is sub-exponential with parameters $(\nu, \alpha)$. Then

\begin{align*}
\Pb[X \geq \mu+t] \leq \begin{cases}e^{-\frac{t^{2}}{2 \nu^{2}}} & \text { if } 0 \leq t \leq \frac{\nu^{2}}{\alpha} \\ e^{-\frac{1}{2 \alpha}} & \text { if } t>\frac{\nu^{2}}{\alpha}\end{cases}
\end{align*}

which can be equivalently written as follows:

\begin{align*}
\Pb[X \geq \mu+t] \leq \exp \left\{-\frac{1}{2} \min \left\{\frac{t}{\alpha}, \frac{t^{2}}{\nu^{2}}\right\}\right\}.
\end{align*}

(Restatement of Theorem 2.2 in [Rigollet and Hütter, 2015]) Assume that the linear model holds where the noise $\varepsilon \sim \operatorname{subG}_{n}\left(\sigma^{2}\right)$. Then the least squares estimator $\hat{\theta}^{\mathrm{LS}}$ satisfies

\begin{align*}
\mathbb{E}\left[\operatorname{MSE}\left(\bX \hat{\theta}^{\mathrm{LS}}\right)\right]=\frac{1}{n} \mathbb{E}\left|\bX \hat{\theta}^{\mathrm{LS}}-\bX \theta^{*}\right|_{2}^{2} \lesssim \sigma^{2} \frac{r}{n}
\end{align*}

where $r=\operatorname{rank}\left(\bX^{\top} \bX\right)$. Moreover, for any $\delta>0$, with probability at least $1-\delta$, it holds

\begin{align*}
\operatorname{MSE}\left(\bX \hat{\theta}^{\mathrm{LS}}\right) \lesssim \sigma^{2} \frac{r+\log (1 / \delta)}{n}
\end{align*}

§.§ Formulation for Design to Reduce MSE

§.§ Loss is convex

The loss function

\begin{align*}
\L_n(\pi, \bb, \bSigma) = \frac{1}{n} \left(\sum_{a,a'}\bw(a)^\top\bA_{\bb,\bSigma}^{-1}\bw(a')\right)
\end{align*}

for any arbitrary design proportion $\bb\in\triangle(\A)$ and co-variance matrix $\bSigma$ is strictly convex.

Let $\bb, \bb' \in \triangle(\A)$, so that $\bA_{\bb}$ and $\bA_{\bb'}$ are invertible. Recall that we have the loss for a design proportion $\bb$ as

\begin{align*}
\L_n(\pi, \bb, \bSigma) &= \frac{1}{n} \left(\sum_{a,a'}\bw(a)^\top\bA_{\bb,\bSigma}^{-1}\bw(a')\right) \overset{(a)}{=} \frac{1}{n} \Tr\left(\sum_{a,a'}\bw(a)^\top\bA_{\bb,\bSigma}^{-1}\bw(a')\right) \\
&= \frac{1}{n}\Tr\left(\bA_{\bb,\bSigma}^{-1}\sum_{a,a'}\bw(a)\bw(a')^\top\right) \\
&= \frac{1}{n}\Tr\left(\bV\bA_{\bb,\bSigma}^{-1}\right)
\end{align*}

where, in $(a)$ we can introduce the trace as the R.H.S. is a scalar quantity,

$\bw(a) = \pi(a)\bx(a)$ and $\bV = \sum_{a,a'}\bw(a)\bw(a')^\top$. Similarly for a $\lambda \in[0,1]$ we have

\begin{align*}
\L_n(\pi, \lambda\bb + (1-\lambda)\bb', \bSigma) = \frac{1}{n} \Tr\left(\bA^{-1}_{\bb,\bb',\bSigma}\sum_{a,a'}\bw(a)\bw(a')^\top\right) = \frac{1}{n} \Tr\left(\bV\bA^{-1}_{\bb,\bb',\bSigma}\right).
\end{align*}

Let the matrix $\bA_{\bb,\bb',\bSigma}$ be defined as

\begin{align*}
\bA_{\bb,\bb',\bSigma} \coloneqq \lambda\bA_{\bb,\bSigma} + (1-\lambda)\bA_{\bb',\bSigma}.
\end{align*}

Now observe that

\begin{align*}
\bA_{\bb,\bb',\bSigma} = \lambda\bA_{\bb,\bSigma} + (1-\lambda)\bA_{\bb',\bSigma}  = \sum_{a=1}^A\left(\lambda\bb(a) + (1-\lambda)\bb'(a)\right)\tx(a)\tx(a)^\top.
\end{align*}

Also observe that this is a positive semi-definite matrix. Now using Lemma 1 from [Whittle, 1958] we can show that

\begin{align*}
\left(\lambda\bA_{\bb,\bSigma} + (1-\lambda)\bA_{\bb',\bSigma}\right)^{-1} \prec \lambda\bA_{\bb,\bSigma}^{-1} + (1-\lambda)\bA_{\bb',\bSigma}^{-1}
\end{align*}

for any positive semi-definite matrices $\bA_{\bb}, \bA_{\bb'}$, and $\lambda\in[0,1]$.
Now taking the trace on both sides we get

\begin{align*}
\Tr\left(\lambda\bA_{\bb,\bSigma} + (1-\lambda)\bA_{\bb',\bSigma}\right)^{-1} \prec \Tr\lambda\bA_{\bb,\bSigma}^{-1} + \Tr(1-\lambda)\bA_{\bb',\bSigma}^{-1}.
\end{align*}

Now using Lemma 2 from Whittle, 1958 we can show that

\begin{align*}
\Tr\left(\lambda\bV\bA_{\bb,\bSigma} + (1-\lambda)\bV\bA_{\bb',\bSigma}\right)^{-1} \prec \Tr\lambda\bV\bA_{\bb,\bSigma}^{-1} + \Tr(1-\lambda)\bV\bA_{\bb',\bSigma}^{-1}.
\end{align*}

for any positive semi-definite matrix $\bV$. This implies that

\begin{align*}
\L_n(\pi, \lambda\bb + (1-\lambda)\bb', \bSigma) < \lambda\L_n(\pi, \bb, \bSigma) + (1-\lambda)\L_n(\pi, \bb', \bSigma).
\end{align*}

Hence, the loss function is convex.

((Bound on variance)

We can use  singular value decomposition of $\bSigma_*$ as $\bSigma_* = \bU \bD \bP^{\top}$ with orthogonal matrices $\bU, \bP^{\top}$ and $\bD=\operatorname{diag}\left(\lambda_{1}, \ldots, \lambda_{d}\right)$ where $\lambda_{i}$ denotes a singular value. Then we can bound $\bx(a)^{\top} \bSigma_* \bx(a)$ as

\begin{align*}
\left\|\bx(a)^{\top} \bSigma_* \bx(a)\right\|=\left\|\bx(a)^{\top} \bU \bD \bP^{\top} \bx(a)\right\|%\\
&\overset{(a)}{=}\left\|\bu^{\top} \bD \mathbf{p}\right\| \leq\left\|\bu^{\top}\right\| \max_{i}|\lambda_{i}| \left\|\mathbf{p} \right\| \\
&\overset{(b)}{=} \| \bx(a)\| \max_{i} |\lambda_{i}| \left\| \bx(a)\right\|
= \max_{i}|\lambda_{i}|\left\| \bx(a) \right\|^{2}
\end{align*}

where in $(a)$ we have $\bu = \bU^{\top} \bx(a)$,  $\mathbf{p} = \bP^{\top} \bx(a)$ and $(b)$ uses the fact that $\left\|\bU^{\top} \bx(a)\right\|=\|\bx(a)\|$ for any orthogonal matrix $\bU^{\top}$. Similarly we can show that $\left\|\bx(a)^{\top} \bSigma_* \bx(a)\right\|\geq \min_{i}|\lambda_{i}|\left\| \bx(a) \right\|^{2}$. Let $H_L^2\leq \|\bx(a)\|^2\leq H_U^2$ for any $a\in[A]$.
This implies that

\begin{align*}
&\underbrace{\min_{i}|\lambda_{i}|H_L^2}_{\sigma^2_{\min}} \leq \min_{i}|\lambda_{i}|\left\| \bx(a) \right\|^{2} \leq \underbrace{\bx(a)^{\top} \bSigma_* \bx(a)}_{\sigma^2(a)} %\\
\leq \max_{i}|\lambda_{i}|\left\| \bx(a) \right\|^{2}\leq \underbrace{\max_{i}|\lambda_{i}|H^2_U}_{\sigma^2_{\max}}
\end{align*}

§.§ Loss Gradient is Bounded

Let $\bb, \bb' \in \triangle(\A)$, so that $\bA_{\bb,\bSigma}$ and $\bA_{\bb',\bSigma}$ are invertible and define $\bV = \sum_{a,a'}\bw(a)\bw(a')^{\top}$. Then the gradient of the loss function is bounded such that

\begin{align*}
\|\nabla_{\bb(a)}\L(\pi,\bb,\bSigma) - \nabla_{\bb(a)}\L(\pi,\bb',\bSigma)\|_2 \leq C_{\kappa}
\end{align*}

where, the

\begin{align*}
C_\kappa = \frac{\lambda_d(\bV)H^2_U}{\sigma^2(a)\left(\min_{a'\in\A}\frac{\bb(a')}{\sigma(a')^2} \lambda_{\min }\left(\sum_{a=1}^A \bw(a) \bw(a)^{\top}\right)\right)^2} \\
+ \frac{\lambda_1(\bV)H^2_U}{\sigma^2(a)\left(\min_{a'\in\A}\frac{\bb'(a')}{\sigma(a')^2} \lambda_{\min }\left(\sum_{a=1}^A \bw(a) \bw(a)^{\top}\right)\right)^2}.
\end{align*}

Let $\bb, \bb' \in \triangle(\A)$, so that $\bA_{\bb,\bSigma}$ and $\bA_{\bb',\bSigma}$ are invertible.
Observe that the gradient of the loss is given by

\begin{align*}
\nabla_{\bb(a)}\L(\pi,\bb,\bSigma) &= \nabla_{\bb(a)} \Tr\left(\sum_{a,a'}\bw(a)^{\top}\bA^{-1}_{\bb,\bSigma}\bw(a')\right) \\
&\overset{(a)}{\leq} \lambda_1(\bV)\nabla_{\bb(a)}\Tr(\bA^{-1}_{\bb,\bSigma}) \\
&\overset{}{=} -\lambda_1(\bV)\Tr\left(\left(\dfrac{\bw(a)\bw(a)^\top}{\sigma^2(a)} \right)\bA^{-2}_{\bb,\bSigma}\right)\\
&= - \lambda_1(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb,\bSigma} \bw(a)\right\|^2_2%\\
% &\leq \lambda_1(\bV)\Tr(\bw(a)\bw(a)^\top B^{-2}_L)\Tr(\bA^{-2}_{\bb,\bSigma})
\end{align*}

where, in $(a)$ we denote $\bV = \sum_{a,a'}\bw(a)\bw(a')^{\top}$.
Similarly the gradient of the loss is lower bounded by

\begin{align*}
&\nabla_{\bb(a)}\L(\pi,\bb,\bSigma) \geq - \lambda_d(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb,\bSigma} \bw(a)\right\|^2_2
% \lambda_d(\bV)\Tr\left(\bw(a)\bw(a)^\top B^{-2}_U\right)\Tr\left(\bA^{-2}_{\bb,\bSigma}\right)
\end{align*}

which yields a bound on the gradient difference as

\begin{align*}
\|&\nabla_{\bb(a)}\L(\pi,\bb,\bSigma) - \nabla_{\bb'(a)}\L(\pi,\bb',\bSigma)\|_2\\
&\leq \left\|\lambda_d(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb,\bSigma} \bw(a)\right\|^2_2 - \lambda_1(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb',\bSigma} \bw(a)\right\|^2_2\right\|_2\\
&\leq \left|\lambda_d(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb,\bSigma} \bw(a)\right\|^2_2\right|+ \left|\lambda_1(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb',\bSigma} \bw(a)\right\|^2_2\right|.%\\
% &\leq \left\|\lambda_d(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb,\bSigma} \bw(a)\right\|^2_2 - \lambda_1(\bV) \dfrac{1}{\sigma^2(a)}\left\|\bA^{-1}_{\bb,\bSigma} \bw(a)\right\|^2_2\right\|_2
\end{align*}

So now we focus on the quantity

\begin{align*}
\left\|\bA^{-1}_{\bb,\bSigma} \bw(a)\right\|^2_2\leq \|\bA^{-1}_{\bb,\bSigma}\|^2_2\|\bw(a)\|^2_2 \leq \|\bA^{-1}_{\bb,\bSigma}\|^2_2 H^2_U.
\end{align*}

Now observe that when $\bb(a)\in\Delta(\A)$ and initialized uniform randomly, then the optimization in (<ref>) results in a non-singular $\bA^{-1}_{\bb,\bSigma}$ if each action has been sampled at least once which is satisfied by . So now we need to bound the minimum eigenvalue of $\bA^{}_{\bb,\bSigma}$ denoted as $\lambda_{\min}(\bA^{}_{\bb,\bSigma})$. Using Lemma 7 of Fontaine et al., 2021 we have that for all $\bb \in \Delta(\A)$,

\begin{align*}
\min_{a \in[A]} \frac{\bb(a)}{\sigma(a)^2} \sum_{a=1}^A \bw(a) \bw(a)^{\top} \preccurlyeq \sum_{a=1}^A \frac{\bb(a)}{\sigma(a)^2} \bw(a) \bw(a)^{\top} .
\end{align*}

And finally

\begin{align*}
\min_{a \in[A]} \frac{\bb(a)}{\sigma(a)^2} \lambda_{\min }\left(\sum_{a=1}^A \bw(a) \bw(a)^{\top}\right) \leq \lambda_{\min}(\bA^{}_{\bb,\bSigma})
\end{align*}

This implies that

\begin{align*}
\lambda_{\min}(\bA^{-1}_{\bb,\bSigma}) \leq \dfrac{1}{\min_{a \in[A]} \frac{\bb(a)}{\sigma(a)^2} \lambda_{\min }\left(\sum_{a=1}^A \bw(a) \bw(a)^{\top}\right)}
\end{align*}

Plugging everything back we get that

\begin{align*}
\|\nabla_{\bb(a)}\L(\pi,\bb,\bSigma) - \nabla_{\bb'(a)}\L(\pi,\bb',\bSigma)\|_2 &\leq \dfrac{\lambda_d(\bV)H^2_U}{\sigma^2(a)\left(\min_{a'\in\A}\frac{\bb(a')}{\sigma(a')^2} \lambda_{\min }\left(\sum_{a=1}^A \bw(a) \bw(a)^{\top}\right)\right)^2} \\
&\qquad + \dfrac{\lambda_1(\bV)H^2_U}{\sigma^2(a)\left(\min_{a'\in\A}\frac{\bb'(a')}{\sigma(a')^2} \lambda_{\min }\left(\sum_{a=1}^A \bw(a) \bw(a)^{\top}\right)\right)^2}.
\end{align*}

The claim of the lemma follows.

§.§ Kiefer-Wolfowitz Equivalence

§.§ Weighted Least Square Estimator that Minimizes MSE

(Weighted Least Square)
Let $\bA_{\bb,\bSigma} = \sum_{a,a'}\lceil \bb(a)n\rceil\tx(a)\tx(a')^{\top}$ be the design matrix such that each action is sampled according to $\bb\in \triangle(A)$. Define the weighted least square estimate $\wtheta_{n}$ by (<ref>).
Then the weighted least square in  (<ref>) estimate minimizes $\E_{\D}\left[\left( \sum_{a=1}^A\bw(a)^{\top}(\wtheta_n-\btheta_*)\right)^{2}\right]$ by minimizing the quantity $\sum_{a,a'} \bw(a)^{\top}\bA_{\bb,\bSigma}^{-1}\bw(a')$.

from <ref> that the Weighted least square estimate is

\begin{align*}
\wtheta_{n} \overset{(a)}{=}\arg\min_{\btheta}\sum_{t=1}^{n}\frac{1}{\sigma^{2}(I_{t})}(R(I_t)-\bv(I_t)^{\top}\btheta)^{2}
& \overset{(b)}{=}\left(\bX_{n}^{\top}\bSigma_{n}^{-1}\tX_{n}\right)^{-1}\bX_{n}^{\top}\bSigma_{n}^{-1}\bR_{n}\\
& =\left(\bX_{n}^{\top}\bSigma_{n}^{-1}\bX_{n}\right)^{-1}\bX_{n}^{\top}\bSigma_{n}^{-1}\left(\bX_{n}\btheta_{*}+\mathbf{\eta}\right)
\end{align*}

where, in $(a)$ the $I_t$ is the action sampled at timestep $t$, and in $(b)$ we define the matrix $\bSigma_n\in\mathbb{R}^{n\times n}$. The  $\mathbf{diag}(\bSigma_n) = [\sigma^2(I_1), \sigma^2(I_2), \ldots, \sigma^2(I_n)]$.
It follows then the value estimate

\begin{align*}
Y_n = \sum_{a=1}^A \bw(a)^{\top}\wtheta_n &= \sum_{a=1}^A \bw(a)^{\top}(\bX_n^{\top}\bSigma_n^{-1}\bX_n)^{-1}\bX_n^{\top}\bSigma_n^{-1}\bR_n \\
&= \sum_{a=1}^A \bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\bX_n^{\top}\bSigma^{-1}_n\left(\bX_n\btheta_* + \mathbf{\eta}\right)\\
&= \sum_{a=1}^A\bw(a)^{\top}\btheta_* + \sum_{a=1}^A\bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\bX_n^{\top}\bSigma^{-1}_n\eta.
\end{align*}

This means that

\begin{align*}
\sum_{a=1}^A \bw(a)^{\top}(\wtheta_n - \btheta_*) &= \sum_{a=1}^A\bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\bX_n^{\top}\bSigma^{-1}_n\eta \\
&= \sum_{a=1}^A\bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1/2}_n\bSigma^{-1/2}_n\bX_n)^{-1}\bX_n^{\top}\bSigma^{-1}_n\eta
\end{align*}

Again, as $\eta$ is bounded so we have  $\eta\sim\SG(0,\bSigma_n)$ where $\SG$ denote the sub-Gaussian distribution. Then we can show that

\begin{align*}
&\left(\sum_{a=1}^A \bw(a)^{\top}(\wtheta_n - \btheta_*)\right)^2 \\
&\sim \SE\left(0, \sum_{a=1}^A\bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\bX_n^{\top}\bSigma^{-1}_n\E[\eta\eta^\top]\bSigma^{-1}_n\bX_n^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\sum_{a=1}^A\bw(a)\right)\\
&\overset{(a)}{\sim} \SE\left(0, \sum_{a=1}^A\bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\bX_n^{\top}\bSigma^{-1}_n\bSigma^{}_n\bSigma^{-1}_n\bX_n^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\sum_{a=1}^A\bw(a)\right)\\
&\sim \SE\left(0, \sum_{a,a'}\bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\bw(a')\right)
\end{align*}

where, $(a)$ follows as $\eta\sim\SG(0,\bSigma_n)$ and $\SE$ denotes sub-exponential distribution. Now using sub-exponential concentration inequality in <Ref>, setting

\nu^2=\sum_{a,a'}\bw(a)^{\top}(\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}\bw(a'),

and $\alpha = \nu$, we can show that

\begin{align*}
\Pb&\left(\left(\sum_{a=1}^A\bw(a)^{\top}(\wtheta_n - \btheta_*)\right)^2 > t \right) \leq \delta, \qquad \text{ if } t \in (0,1]\\
\Pb&\left(\left(\sum_{a=1}^A\bw(a)^{\top}(\wtheta_n - \btheta_*)\right)^2 > t^2 \right) \leq \delta, \qquad \text{ if } t > 1 .
\end{align*}

Combining the above two we can show that

\begin{align*}
\Pb&\left(\left(\sum_{a=1}^A\bw(a)^{\top}(\wtheta_n - \btheta_*)\right)^2 > \min\{t, t^2\} \right) \leq \delta, \forall t > 0.
\end{align*}

Further define matrix $\bbSigma_n \in \mathbb{R}^{d\times d}$ as $\bbSigma_n^{-1} \coloneqq (\bX_n^{\top}\bSigma^{-1}_n\bX_n)^{-1}$.

Hence using sub-exponential concentration inequality we can show that

\begin{align*}
\Pb\left(\left(\sum_{a=1}^A\bw(a)^{\top}(\wtheta_n - \btheta_*)\right)^2 > \min\left\{\sqrt{2\sum_{a,a'}\bw(a)^{\top}\bbSigma_n^{-1}\bw(a') \log (1 / \delta)}, 2\sum_{a,a'}\bw(a)^{\top}\bbSigma_n^{-1}\bw(a') \log (1 / \delta)\right\} \right) \leq \delta.
\end{align*}

This means that we have with probability $(1-\delta)$ that

\begin{align*}
\left(\sum_{a=1}^A\bw(a)^{\top}(\wtheta_n - \btheta_*)\right)^2 &\leq \min\left\{\sqrt{2\sum_{a,a'} \bw(a)^{\top}\bbSigma^{-1}\bw(a') \log (1 / \delta)}, 2\sum_{a,a'}\bw(a)^{\top}\bbSigma_n^{-1}\bw(a') \log (1 / \delta)\right\}\\
&= 2\sum_{a,a'} \bw(a)^{\top}\bbSigma_n^{-1}\bw(a') \log (1 / \delta).
\end{align*}

Recall that we have sampled each action till $n$ in some proportion $\bb\in\triangle(\A)$. Then we have that

\begin{align*}
\bbSigma_n = \bX_n^{\top}\bSigma^{-1}_n\bX_n &= \sum_{a=1}^A\lceil \bb(a)n\rceil\pi^2(a)\sigma^{-2}(a)\bv(a)\bv(a)^{\top}\\
&= \sum_{a=1}^A\lceil \bb(a)n\rceil\tx(a)\tx(a')^{\top} = n\bA_{\bb,\bSigma}.
%\qquad \bX_n^{\top}\bSigma_n\bX_n = \sum_{a=1}^A\lceil b(a)n\rceil\sigma^2(a)\bv(a)\bv(a)^{\top} = n\bbA_b.
\end{align*}

It follows then that

\begin{align*}
\left(\sum_{a=1}^A\bw(a)^{\top}(\wtheta_n - \btheta_*)\right)^2 &\leq \frac{2}{n}\sum_{a,a'} \bw(a)^{\top}\bA_{\bb,\bSigma}^{-1}\bw(a') \log (1 / \delta).
\end{align*}

Hence, using the weighted least square estimate we can show that it

minimizes $\E_{\D}\left[\left( \sum_{a=1}^A\bw(a)^{\top}(\wtheta_n-\btheta_*)\right)^{2}\right]$, by minimizing the quantity $\sum_{a,a'} \bw(a)^{\top}\bA_{\bb,\bSigma}^{-1}\bw(a')$ where $\bA_{\bb,\bSigma} = \sum_{a,a'}\lceil \bb(a)n\rceil\tx(a)\tx(a')^{\top}$.

Note that the estimator $\wtheta_n$ is an unbiased estimator of $\btheta_*$. This can be shown as follows

\begin{align*}
\E\left[\wtheta_n\right] - \btheta_* &= \E\left[\left(\bX_{n}^{\top}\bSigma_{n}^{-1}\bX_{n}\right)^{-1}\bX_{n}^{\top}\bSigma_{n}^{-1}\bR_{n}\right] - \btheta_*\\
&= \E\left[\left(\bX_{n}^{\top}\bSigma_{n}^{-1}\bX_{n}\right)^{-1}\bX_{n}^{\top}\bSigma_{n}^{-1}\left(\bX_{n}\btheta^{*}+\mathbf{\eta}\right)\right] - \btheta_*\\
&=  \E\left[\left(\bX_{n}^{\top}\bSigma_{n}^{-1}\bX_{n}\right)^{-1}\bX_{n}^{\top}\bSigma_{n}^{-1}\bX_{n}\btheta^{*}\right]+\E\left[\left(\bX_{n}^{\top}\bSigma_{n}^{-1}\bX_{n}\right)^{-1}\bX_{n}^{\top}\bSigma_{n}^{-1}\mathbf{\eta}\right] - \btheta_*\\
&= \btheta_* + \left(\bX_{n}^{\top}\bSigma_{n}^{-1}\bX_{n}\right)^{-1}\bX_{n}^{\top}\bSigma_{n}^{-1}\E\left[\mathbf{\eta}\right] - \btheta_* \overset{(a)}{=} 0
\end{align*}

where, $(a)$ follows as noise is zero mean.

§ BANDIT REGRET PROOFS

§ REGRET LOWER BOUND

§ ADDITIONAL EXPERIMENTS

§ TABLE OF NOTATIONS

Notations     Definition

$\pi(a)$      Target policy probability for action $a$

$\bb(a)$      Behavior policy probability for action $a$

$\bx(a)$      Feature of action $a$

$\btheta_*$      Optimal mean parameter

$\wtheta_n$      Estimate of $\btheta_*$

$\mu(a) = \bx^\top\btheta_*$      Mean of action $a$

$\wmu_t(a) = \bx^\top\wtheta_t$      Empirical mean of action $a$ at time $t$

$R_t(a)$      Reward for action $a$ at time $t$

$\bSigma_*$      Optimal co-variance matrix

$\wSigma_t$      Empirical co-variance matrix at time $t$

$\sigma^2(a) = \bx(a)^\top\bSigma_*\bx(a)$      Variance of action $a$

$\wsigma_t^2(a) = \bx(a)^\top\wSigma_t\bx(a)$      Empirical variance of action $a$ at time $t$

$n$     Total budget

$T_n(a)$      Total Samples of action $a$ after $n$ timesteps

Table of Notations