EconBase
← Back to paper

Adaptive Experimental Design for Policy Learning

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

38,839 characters · 15 sections · 29 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

Adaptive Experimental Design for Policy Learning

abstractThis study investigates the contextual best arm identification (BAI) problem, aiming to design an adaptive experiment to identify the best treatment arm conditioned on contextual information (covariates). We consider a decision-maker who assigns treatment arms to experimental units during an experiment and recommends the estimated best treatment arm based on the contexts at the end of the experiment. The decision-maker uses a policy for recommendations, which is a function that provides the estimated best treatment arm given the contexts. In our evaluation, we focus on the worst-case expected regret, a relative measure between the expected outcomes of an optimal policy and our proposed policy. We derive a lower bound for the expected simple regret and then propose a strategy called Adaptive Sampling-Policy Learning (PLAS). We prove that this strategy is minimax rate-optimal in the sense that its leading factor in the regret upper bound matches the lower bound as the number of experimental units increases.

Introduction

In this study, we design an adaptive experiment for policy learning. We consider the problem of decision-making given multiple treatment arms, such as arms in slot machines, diverse therapies, and distinct unemployment assistance programs. The primary objective is to identify the best treatment arm for individuals given covariates, often referred to as context, at the end of an experiment. For this purpose, we aim to learn a policy that recommends the conditional best treatment arm by using data adaptively collected via an experiment.

In our setting, at each round of an adaptive experiment, a decision-maker sequentially observes a context (covariate) and assigns one of the treatment arms to a research subject based on past observations and the observed contexts. At the end of the experiment, the decision-maker recommends an estimated best treatment arm conditional on a context.

We design the adaptive experiment by developing a strategy that the decision-maker follows. A strategy is defined as a pair of a sampling rule and a recommendation rule. In the adaptive experiment, the decision-maker assigns treatment arms following the sampling rule during the experiment and recommends a treatment arm following the recommendation rule at the end of the experiment.

We measure the performance of a strategy using the expected simple regret, which is the difference between the maximum expected outcome that could be achieved with full knowledge of the distributions of the treatment arms and the expected outcome of the treatment arm recommended by the decision-maker's strategy. Our goal is to develop a strategy that minimizes the expected simple regret.

The challenge in our problem arises from the need for context-specific recommendations. Unlike prior studies that do not consider contextual information, developing a model that captures the relationship between context and outcomes becomes imperative.

To address this issue, we define a function suggesting a treatment arm given a context as a policy. By using a policy, we restrict strategies to ones with recommendation rules using a policy learned from observations obtained from an adaptive experiment. Policy learning has been extensively studied in causal inference and reinforcement learning dudik2011doubly,Swaminathan2015,Kitagawa2018,AtheySusan2017EPL,ZhouZhengyuan2018OMPL, but to the best of our knowledge, adaptive experimental design for policy learning has not been fully explored. Note that there are existing studies that address contextual BAI aiming to find the best treatment arm marginalized over the contextual distribution Russac2021,Kato2021Role,simchi2024experimentation, motivated by the studies of efficient average treatment effect estimation Laan2008TheCA,Hahn2011,Kato2020adaptive,cook2023semiparametric.

Our problem corresponds to a generalization of best arm identification Bubeck2009,Bubeck2011,Audibert2010, an instance of the stochastic multi-armed bandit (MAB) problem Thompson1933,Lai1985. Therefore, we refer to the problem as contextual fixed-budget BAI, as well as an adaptive experimental design for policy learning.

\paragraph{Contribution.} We propose a strategy that assigns treatment arms following Adaptive Sampling (AS) and recommends a treatment arm using Policy Learning (PL). In the AS rule, a decision-maker assigns a treatment arm to an experimental unit based on a probability depending on the variances of the experimental units' outcomes. Because the variances are unknown, the decision-maker estimates them during the experiment and continues updating the assignment probability. At the end of the experiment, the decision-maker trains a policy using observations obtained in the experiment and recommends a treatment arm using the trained policy. We refer to our strategy as the PLAS strategy.

To design an optimal strategy, we first develop a lower bound (theoretical limit) for the expected simple regret. Subsequently, we design a strategy and evaluate its upper bound (performance) by comparing it to the lower bound.

In the evaluation, given the inherent uncertainties, we use the minimax criterion for performance assessment, which evaluates the worst-case scenario among a set of distributions. The minimax approach has garnered attention in studies about experimental design, including BAI Bubeck2009,Bubeck2011,Carpentier2016,Ariu2021,yang2022minimax,Komiyama2022. A critical quantity in the minimax evaluation is the gap between the expected outcomes of the best treatment and the other suboptimal treatment arms, referred to as the average treatment effects in the literature of causal inference. The worst-case distributions are characterized by gaps approaching zero at a rate of order $\sqrt{T}$ Bubeck2009,Bubeck2011, where $T$ denotes the sample size (total rounds of an adaptive experiment).

Our research identifies the leading factor in the lower bound as the variances of potential outcomes, also providing a variance-dependent sampling rule. We subsequently show that the PLAS strategy is asymptotically minimax optimal, as its foremost factor of the worst-case expected simple regret aligns with the lower bound.

In summary, our contributions include: (i) a lower bound for the worst-case expected simple regret; (ii) the PLAS strategy with a closed-form target assignment ratio, characterized by the variances of outcomes; and (iii) the asymptotic minimax optimality of the PLAS strategy. These findings contribute to a variety of subjects, including decision theory and causal inference, in addition to BAI.

\paragraph{Organization.} The structure of this paper is as follows: Section (ref) defines our problem. Section (ref) develops lower bounds for the worst-case expected simple regret. Section (ref) introduces the PLAS strategy, and Section (ref) presents upper bounds for the proposed strategy and its asymptotic minimax optimality. Further related work is introduced in Appendix (ref).

Problem Setting

This study considers an adaptive experiment with a fixed budget (sample size) $T\in\mathbb{N}$, a set of treatment arms $[K] \coloneqq \{1, 2, \dots, K\}$, and a decision-maker who aims to identify the context-conditional best treatment arm. In each round $t \in [T]$, experimental units sequentially visit, and the decision-maker can assign treatment arms to them. At the end of the experiment, the decision-maker recommends an estimated context-conditional best treatment arm.

Potential Outcomes

Following the Neyman-Rubin causal model Neyman1923,Rubin1974, let $Y^a\in\mathcal{Y}$ be a potential random outcome of treatment arm $a\in [K]$, where $\mathcal{Y}\subset \mathbb{R}$ denotes a set of possible outcomes. We also define $X \in \mathcal{X}$ as a context, also called covariates, that characterizes an experimental unit, where $\mathcal{X}\subset \mathbb{R}^d$ denotes a set of possible $d$-dimensional covariates. We define a tuple $(Y^1, \dots, Y^K, X)$ as a $(\mathcal{Y}^K \times \mathcal{X})$-valued random variable following a probability distribution $P \in \mathcal{M}(\mathcal{Y}^K \times \mathcal{X})$, where $\mathcal{M}(\mathcal{T})$ denotes the set of all Borel probability measures on a topological space $\mathcal{T}$. Let $\mathbb{E}_P$ and $\mathrm{Var}_{P}$ be the expectation and variance operators under $P$, respectively. For each $P \in\mathcal{P}$, let us denote the first and second moments of the potential outcome as $\mu^a(P)(x) \coloneqq \mathbb{E}_P[Y^a \mid X=x]$ and $\nu^a(P)(x) \coloneqq \mathbb{E}_{P}[(Y^a)^2\mid X=x]$, respectively.

In our setting, a distribution $P$ of $(Y^1,\dots, Y^K, X)$ belongs to a bandit model $\mathcal{P}$ defined below. We refer to the bandit model as a location-shift bandit model.

definition[Location-shift bandit model] Let $(\sigma^a)^2 \colon \mathcal{X} \to (0, +\infty)$ be a function that is exogenously given but unknown to the decision-maker. Let $\zeta$ be a distribution of $X$ that is exogenously given but unknown to the decision-maker, which is supported on $\mathcal{X}$. Then, a location-shift bandit model $\mathcal{P}$ is defined as follows: \begin{align} &\mathcal{P} \coloneqq \mathcal{P}_\zeta \coloneqq \Big\{P \in \mathcal{M}(\mathcal{Y}^K \times \mathcal{X}) \colon \forall a \in [K]\ \ \ \forall x \in \mathcal{X},\ \ \ \mathrm{Var}_P(Y^a \mid X=x) = \sigma^a(x), \\ &\ Y^a - \mu^a(P)(x)\ \mathrm{is}\ \mathrm{zeo}\mathchar`-\mathrm{mean}\ \mathrm{subgaussian}\ \mathrm{given}\ x,\ \ \ \mu^a(P)(x) \in (-\infty, +\infty),\ \ \ \mathrm{marg}_{\mathcal{X}}(P) = \zeta \Big\}, \end{align} where $\mathrm{marg}_{\mathcal{X}}(P)$ denotes the marginal distribution of $P$ on $\mathcal{X}$.

Our lower and upper bounds are characterized by $\sigma^a(x)$, given independently of the experiment. Location-shift models are a commonly employed assumption in statistical analysis LehmCase98. A key example is a Gaussian distribution, where for all $P$, the variances are fixed and only mean parameters shift. Note that we can omit the condition $|\nu^a(P)(x)| < \overline{C}$ from the boundedness of the variance. However, we introduce it to simplify the definition of our strategies.

Adaptive Experiment

We consider a decision-maker who aims to identify the best treatment arm \[a^*(P)(x) \in \operatorname*{arg\,max}_{a\in [K]}\mu^a(P)(x)\] for each context $x \in \mathcal{X}$ through an adaptive experiment. A fixed number of rounds $T$, called the budget or sample size, is exogenously given. At each round $t \in [T],$ the decision-maker uses the following procedure:

enumerate[topsep=0pt, itemsep=0pt, partopsep=0pt, leftmargin=*] • A potential outcome $(Y^1_t, Y^2_t, \dots, Y^K_t, X_t)$ is generated from $P$. • The decision-maker observes a context $X_t$. • The decision-maker assigns treatment arm $A_t$ to an experimental unit based on past observations $(X_s, A_s, Y_s)^{t-1}_{s=1}$ and the context $X_t$. • The decision-maker observes an outcome of the assigned treatment arm, $Y_t = \sum_{a\in[K]}\mathbbm{1}[A_t = a]Y^a_t$.

This process is referred to as the exploration phase. Note that outcomes from unassigned treatment arms remain unobserved. This setting is called the bandit feedback or Rubin causal model Neyman1923,Rubin1974. Through the exploration phase, the decision-maker obtains observations $(X_t, A_t, Y_t)^T_{t=1}$. After round $T$, the decision-maker recommends an estimated best treatment arm $\widehat{a}_T(x)$ for each context $x$ given the observations $(X_t, A_t, Y_t)^T_{t=1}$.

Strategy with Policy Learning

A strategy of the decision-maker defines which treatment arm to assign in each period during the exploration phase and which treatment arm to recommend as an estimated best arm for each context. A strategy is defined as a pair $((A_t)_{t\in[T]}, \widehat{a}_T)$ in which:

itemize[topsep=0pt, itemsep=0pt, partopsep=0pt, leftmargin=*] • the sampling rule $(A_t)_{t\in[T]}$ determines which treatment arm $A_t$ to assign in round $t$ based on the past observations. In other words, $A_t$ is $\mathcal{G}_t$-measurable, where $\mathcal{G}_t \coloneqq \sigma(X_1,A_1,Y_1, \dots, X_{t-1}, A_{t-1}, Y_{t-1}, X_t)$ for each $t \in [T]$. • the recommendation rule $\widehat{a}_T:\mathcal{X}\to [K]$ returns an estimated best treatment arm for each context $\widehat a_T$ based on the observations collected during the exploration phase. In other words, for each $x \in \mathcal{X}$, $\widehat a_T(x)$ is $\mathcal{F}_T$-measurable, where $\mathcal{F}_T \coloneqq \sigma(X_1, A_1, Y_1, \dots, X_T, A_T, Y_T)$.

\paragraph{Policy.} In this study, we impose a restriction on a class of recommendation rules that the decision-maker can use. We assume that there is an exogenously given class of policies $\Pi$, whose typical element is a policy $\pi \colon [K] \times \mathcal{X} \to [0,1]$ that is measurable and satisfies $\sum_a \pi(a, x) = 1$ for each $x$. Here, $\pi(a,x)$ denotes a probability that the decision-maker recommends treatment arm $a \in [K]$ as an estimated best treatment arm for context $x$. With a slight abuse of notation, we write $\pi(a \mid x)$ instead of $\pi(a,x)$. We require the decision-maker to obtain an estimator $\widehat{a}_T(x)$ of $a^*(P)(x)$ as follows: first, the decision-maker constructs a policy $\pi \in \Pi$ based on the observations collected during the exploration phase; then, $\widehat{a}_T(x)$ is drawn from $\pi(\cdot \mid x)$ for each $x$.

\paragraph{Optimal policy.} We evaluate the performance of policies via a simple regret. The value of policy $\pi\in\Pi$ under $P$ is the expected outcome when the decision-maker uses a policy $\pi$, which is defined as \[ Q(P)(\pi) \coloneqq \mathbb{E}_P\left[\sum_{a\in[K]}\pi(a\mid X)\mu^{a}(P)(X); \pi\right], \] and the optimal policy within class $\Pi$ is defined as $ \pi^*(P) \coloneqq \operatorname*{arg\,max}_{\pi\in\Pi} Q(P)(\pi)$.

Given a strategy with a policy $\widehat{\pi}\in\Pi$ of the decision-maker, we define a simple regret for each context $x \in \mathcal{X}$ under $P \in \mathcal{P}$ as

align*[align* omitted — 142 chars of source]

and the marginalized simple regret is defined as

align*[align* omitted — 192 chars of source]

Then, we define the expected simple regret as $\mathbb{E}_P\left[R_T(P)\left(\widehat{\pi}\right)\right]$, where the expectation is taken over $\widehat{\pi}$. This expected simple regret is our performance measure of interest. We also refer to the expected simple regret $\mathbb{E}_P\left[R_T(P)\left(\widehat{\pi}\right)\right]$ as the policy regret. The decision-maker aims to identify the best treatment arm with a smaller expected simple regret.

\paragraph{Notation.} Let $o(g(x))$ be Landau's notation, and $f(x) = o(g(x))$ implies that $\forall \varepsilon > 0\ \exists x_0\ \forall x > x_0\colon |f(x)| < \varepsilon g(x)$ holds. Let let $\mathrm{thre}(A, a, b) \coloneqq \max\{ \min\{A,a\},b\}$ be a truncation function.

Regret Lower Bound

This section presents a lower bound on the expected simple regret $\mathbb{E}_P\left[R_T(P)\left(\widehat{\pi}\right)\right]$. The lower bound is provided under weak conditions on the policy. Not only does the lower bound offer insights into the difficulty of the problem, but it also helps argue which sample allocations are optimal.

Restriction and Complexity of a Policy Class

To establish lower bounds, we introduce a moderate precondition related to the strategy space of the decision-maker. Specifically, we require that, in the limit, strategies choose all the arms with an equal probability when, for a given covariate $x$, the expected outcomes associated with all arms are identical. Strategies adhering to this criterion are termed null consistent strategies.

definition[Null consistent strategy] We say a strategy is null consistent if the following condition is satisfied: If $\mu^1(P)(x) = \mu^2(P)(x) = \cdots = \mu^K(P)(x)$, then for any $a, b \in [K]$, we have \[ \Big|\mathbb{P}_{P}\left(\widehat{a}_T(X) = a \mid X = x\right) - \mathbb{P}_{P}\left(\widehat{a}_T(X) = b \mid X = x\right)\Big| \to 0 \quad (T\to\infty). \]

Under any null consistent strategies, $\big|\mathbb{P}_{P}\left(\widehat{a}_T(X) = a \mid X = x\right) - 1/K\big| = o(1)$ holds for each $a\in[K]$ as $T\to\infty$ if $\mu^1(P)(x) = \mu^2(P)(x) = \cdots = \mu^K(P)(x)$.

Next, we introduce the Natarajan dimension, a metric that measures the complexity of a policy class $\Pi$ Natarajan1989. Our lower bounds are characterized by the Natarajan dimension.

definition[Natarajan dimension] We say that $\Pi$ shatters $M$ points $\{s_1, s_2, \dots, s_M\} \subseteq \mathcal{X}$ if there exist $f_1, f_{-1}: \{s_1, s_2, \dots, s_M\}\to [K]$ such that \begin{enumerate}[topsep=0pt, itemsep=0pt, partopsep=0pt, leftmargin=*] • for any $j \in [M]$, $f_{-1}(s_j) \neq f_1(s_j)$ holds; • for any $\bm{\sigma} \coloneqq \{\sigma_1, \sigma_2, \dots, \sigma_M\}\in \{\pm 1\}^M$, there exists a policy $\pi\in\Pi$ such that for any $j\in[M]$, it holds that $\pi(s_j) = \begin{cases} f_1(s_j) & \mathrm{if}\ \ \ \sigma_j = 1\\ f_{-1}(s_j) & \mathrm{if}\ \ \ \sigma_j = -1 \end{cases}$. \end{enumerate} The Natarajan dimension of $\Pi$, denoted by $d_{\mathrm{N}}(\Pi)$, is the maximum cardinality of a set shattered by $\Pi$.

Let $d_{\mathrm{VC}}(\Pi)$ be the Vapnik-Chervonenkis (VC) dimension of $\Pi$. Note that when $K = 2$, the Natarajan dimension is equivalent to the VC dimension; that is, when $K = 2$, $d_{\mathrm{VC}}(\Pi) = d_{\mathrm{N}}(\Pi)$ holds.

Regret Lower Bounds

We derive the following lower bounds of the expected simple regret, which hold for any null consistent strategies and depend on the complexity of a policy class $\Pi$, measured by the Natarajan dimension. The proof is shown in Appendix (ref).

theoremThere exists a distribution $\zeta$ on $\mathcal{X}$ such that for any $K \geq 2$, any null consistent strategy $\pi$ with a policy class $\Pi$ such that $d_{\mathrm{N}}(\Pi) = M$ satisfies \begin{align*} &\sup_{P\in\mathcal{P}_\zeta} \sqrt{T}\mathbb{E}_P\left[R(P)\left(\pi\right)\right] \geq \frac{1}{8}\mathop{\mathbb{E}}_{X \sim \zeta} \left[\sqrt{M\sum_{a\in[K]}\left(\sigma^a(X)\right)^2}\right] + o(1)\quad \mathrm{as}\ \ T\to \infty, \end{align*} where $\mathop{\mathbb{E}}_{X \sim \zeta}$ denotes the expectation of a random variable $X$ under the probability distribution $\zeta$.

Our lower bounds also depend on the variances of outcomes $Y^a$.

When $K = 2$, we can obtain a tighter lower bound than the one in Theorem (ref). The proof is provided in Section (ref).

theoremThere exists a distribution $\zeta$ on $\mathcal{X}$ such that for $K = 2$, any null consistent strategy $\pi$ with policy class $\Pi$ such that $d_{\mathrm{VC}}(\Pi) = M$ satisfies \begin{align*} &\sup_{P\in\mathcal{P}_\zeta} \sqrt{T}\mathbb{E}_P\left[R(P)\left(\pi\right)\right] \geq \frac{1}{8}\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{M\left(\sigma^{1}(X) + \sigma^{2}(X)\right)^2}\right] + o(1)\quad \mathrm{as}\ \ T\to \infty. \end{align*}

Here, note that $\sum_{a\in\{1,2\}}\left(\sigma^a(x)\right)^2 \leq \left(\sigma^{1}(x) + \sigma^{2}(x)\right)^2$ holds for each $x\in\mathcal{X}$. Therefore, when $K = 2$, we use the lower bound in Theorem (ref) and when $K\geq 3$, we use the one in Theorem (ref).

The PLAS Strategy

Our strategy consists of the following sampling and recommendation rules. First, we define a target assignment ratio, which is an ideal treatment assignment probability. At each round, $t = 1,2,\dots, T$, our sampling rule randomly assigns a unit to a treatment arm with a probability identical to an estimated target assignment ratio. After the final round $T$, our recommendation rule recommends a treatment arm with the highest value of a policy trained by maximizing empirical policy value. We refer to our strategy as the PLAS strategy. Our strategy depends on hyperparameters $\overline{C} \in (0, \infty)$, which are introduced for technical purpose in the proof and can be set as sufficiently large values. We show a pseudo-code in Algorithm (ref).

Optimal Target Assignment Ratio

We first define a target assignment ratio. The target assignment ratio is the expected value of the sample average of $A_t$ of a strategy ($\frac{1}{T}\sum^T_{t=1}\mathbb{E}_P\left[\mathbbm{1}[A_t = a]\mid X_t = x\right]$) under which a leading factor of its expected simple regret aligns with that of our derived lower bound.

definition[Target assignment ratio] When $K=2$, for each $a\in[K] = \{1, 2\}$, we define the target assignment ratio $w^*$ as \begin{align} w^*(a\mid x) = \frac{\sigma^{a}(x)}{\sigma^{1}(x) + \sigma^{2}(x)}. \end{align} When $K \geq 3$, for each $a\in[K]$, we define the target assignment ratio as \begin{align} w^*(a\mid x) = \frac{ \left(\sigma^{a}(x)\right)^2}{\sum_{b\in[K]} \left(\sigma^{b}(x)\right)^2}. \end{align}

This target assignment ratio is given in the course of proving Theorem (ref), in which we solve $\min_{\bm{w}\in \mathcal{W}} \left\{ \sum_{s\in\mathcal{S}} \max_{a \in [K]} \left\{\sqrt{\frac{\left(\sigma^{a}(s)\right)^2}{w(a\mid s)}}\right\} \right\}$, where $\mathcal{S} \coloneqq \{s_1, s_2, \dots, s_M\} \subseteq \mathcal{X}$, and $\mathcal{W}$ is the set of all measurable functions $w:\mathcal{X}\times [K] \to (0, 1)$ such that $\sum_{a\in[K]} w(a\mid x) = 1$ for each $x\in \mathcal{X}$. The solutions $w^*$, whose explicit forms appear in (ref) and (ref), work as a conjecture for $\frac{1}{T}\sum^T_{t=1}\mathbb{E}_P\left[\mathbbm{1}[A_t = a]\mid X_t = x\right]$.

These target assignment ratios are ex-ante unknown to the decision-maker since the variance $(\sigma^a(x))^2$ is unknown. Therefore, the decision-maker needs to estimate it during an adaptive experiment and employ the estimator as a probability of assigning a treatment arm.

remark[Efficiency gain] For each $a\in[K]$, let $\left(\sigma^a\right)^2$ be the unconditional variance of $Y^a_t$, and $w:[K]\to[0, 1]$ be an assignment ratio such that $\sum_{a\in[K]}w(a) = 1$ when we cannot utilize the contextual information. Then, the following inequality holds: \begin{align*} &\sup_{P\in\mathcal{P}} \sqrt{T}\mathbb{E}_P\left[R(P)\left(\widehat{\pi}\right)\right] \geq \frac{1}{8}\inf_{w\in\mathcal{W}}\max_{a\in[K]}\sqrt{M{\left(\sigma^{a}\right)^2}/{w(a)}}\\ &\geq \frac{1}{8}\inf_{w\in\mathcal{W}}\max_{a\in[K]}\sqrt{M\mathop{\mathbb{E}}_{X \sim \zeta}\left[{\left(\sigma^{a}(X)\right)^2}/{w(a\mid X)}\right]} \geq \frac{1}{8}\inf_{w\in\mathcal{W}}\max_{a\in[K]}\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{M{\left(\sigma^{a}(X)\right)^2}/{w(a\mid X)}}\right]. \end{align*} This result implies that we can minimize a lower bound by using contextual information; that is, strategies utilizing contextual information are more efficient than ones not utilizing contextual information.

Sampling Rule with Adaptive Sampling (AS)

In this section, we describe our sampling rule, referred to as athe AS rule.

In round $t \leq K$, the strategy chooses $A_t = t$, i.e., each arm is pulled once as initialization. In round $t > K$, given an estimated target assignment ratio $\widehat{w}_{t}(a\mid x)$, we assign treatment arm $a$ with probability $\widehat{w}_{t}(a\mid X_t)$. Below, we describe the construction of $\widehat{w}_{t}(a\mid x)$.

In each round $t$, we estimate $w^*$ using the past observations. We first construct estimators $\widehat{\mu}^a_{t}(x)$ and $\widehat{\nu}^a_{t}(x)$ of the first moment $\mu^a(P)(x)$ and the second moment $\nu^a(P)(x)$ of $Y^a$. The estimators constructed to converges to the true functions with probability one, as stated in Assumption (ref), and their absolute values are bounded by $\overline{C}$. Then, given these estimators, we estimate the variances as $\left(\widehat{\sigma}^a_{t}(x)\right)^2 = \mathrm{thre}\Big((\widehat{\sigma}^{\dagger a}_t(x))^2, \overline{C}, 1/\overline{C}\Big)$, where $(\widehat{\sigma}^{\dagger a}_t(x))^2 = \widehat{\nu}^a_{t}(x) - \left(\widehat{\mu}^a_{t}(x)\right)^2$. Lastly, we construct the estimator of the target assignment ratio $\widehat{w}_t(a\mid x)$ by replacing $\sigma^a(x)$ by the estimator $\widehat{\sigma}^a_{t}(x)$.

For obtaining estimators $\widehat{\mu}^a_{t}(x)$, $\widehat{\nu}^a_{t}(x)$, and $\widehat{\sigma}^a_{t}(x)$, we can use nonparametric estimators, such as the nearest neighbor regression estimator and kernel regression estimator, which have been proven to converge to the true function almost surely under a bounded sampling probability $\widehat{w}_t$ by yang2002 and qian2016kernel. It should be noted that we do not assume specific convergence rates for estimators for $\mu^a(P)(x)$ and $w^*$ as the asymptotic optimality of the AIPW estimator can be demonstrated without them Laan2008TheCA,Kato2020adaptive,Kato2021adr.

Recommendation Rule with Policy Learning

The following part presents our recommendation rule. To recommend the conditionally best treatment arm $a^*(P)(x)$, we train a policy $\pi:[K]\times \mathcal{X} \to [0,1]$ by maximizing the empirically approximated policy value function, which we will describe below.

At the end of an experiment, we estimate the policy value $Q(\pi)$ by using the augmented doubly robust estimator, which is defined as follows:

align[align omitted — 141 chars of source]

where

align[align omitted — 249 chars of source]

and $U_T$ is a positive value that approaches infinity as $T \to \infty$. Then, we train a policy as

align[align omitted — 111 chars of source]

By using this trained policy, given $x\in\mathcal{X}$, we recommend $\widehat{a}_T(x) \in [K]$ as the best treatment arm with probability $\widehat{\pi}^{\mathrm{PLAS}}_T\left(\widehat{a}_T(x)\mid x\right)$.

The AIPW estimator debiases the sample selection bias resulting from treatment assignment based on contextual information. Additionally, the AIPW estimator possesses the following properties: (i) its components $\{\widehat{\Gamma}^a_t - \mu^a(P)(x)\}^T_{t=1}$ are a martingale difference sequence, allowing us to employ the martingale limit theorems in derivation of the upper bound; (ii) it has the minimal asymptotic variance among the possible estimators. For example, other estimators with a martingale property, such as the inverse probability weighting (IPW) estimator, may be employed, yet their asymptotic variance would be greater than that of the AIPW estimator. The $t$-th element of the sum in the AIPW estimator utilizes nuisance parameters ($\mu^a(P)(x)$ and $w^*$) estimated from past observations up to round $t-1$ for constructing a martingale difference sequence Laan2008TheCA,hadad2019,Kato2020adaptive,Kato2021adr. For those reasons, this estimator is often used in the context of adaptive experimental design.

algorithm[algorithm omitted — 1,107 chars of source]

Regret Upper Bound

This section provides upper bounds for the expected simple regret of the PLAS strategy. First, we assume the following convergence rate for estimators of $\mu^a(P)(x)$ and $w^*(a\mid x)$.

assumptionFor any $\zeta$, any $P\in\mathcal{P}_\zeta$, and all $a \in [K]$, it holds that \begin{align} &\sup_{x\in\mathcal{X}}\Big|\widehat{w}_T(a\mid x) - w^*(a\mid x) \Big|\xrightarrow{\mathrm{a.s.}} 0,\quad \sup_{x\in\mathcal{X}}\Big|\widehat{\mu}^a_T(x) - \mu^a(P)(x) \Big|\xrightarrow{\mathrm{a.s.}} 0\qquad \mathrm{as}\ \ T\to \infty. \end{align}

Next, we define an entropy integral of a policy class.

definitionGiven the feature domain $\mathcal{X}$, a policy class $\Pi$, a set of $n$ points $\{x_1, \dots, x_n\} \subset \mathcal{X}$, define: \begin{enumerate} • Hamming distance between any two policies $\pi_1$ and $\pi_2$ in $\Pi$: $H(\pi_1, \pi_2) = \frac{1}{n} \sum_{j=1}^n \mathbbm{1}\left[\pi_1(x_j) \neq \pi_2(x_j)\right]$. • $\epsilon$-Hamming covering number of the set $\{x_1, \dots, x_n\}$: $\mathbb{N}_H(\epsilon, \Pi, \{x_1, \dots, x_n\})$ is the smallest number $K$ of policies $\{\pi_1, \dots, \pi_K\}$ in $\Pi$, such that $\forall \pi \in \Pi, \exists \pi_i, H(\pi, \pi_i) \le \epsilon$. • $\epsilon$-Hamming covering number of $\Pi$: $\mathbb{N}_H(\epsilon, \Pi) = \sup\{\mathbb{N}_H(\epsilon, \Pi, \{x_1, \dots, x_m\}) \mid m \ge 1, x_1, \dots, x_m \in \mathcal{X}\}.$ • Entropy integral: $\kappa(\Pi) = \int_0^1 \sqrt{\log \mathbb{N}_H(\epsilon^2, \Pi)}d\epsilon$. \end{enumerate}

The entropy represents the complexity of a policy class, as well as the Natarajan dimension. Between the entropy integral $\kappa(\Pi)$ and the Natarajan dimension $d_{\mathrm{N}}(\Pi)$, $\kappa(\Pi) \leq C\sqrt{\log(d)d_{\mathrm{N}}(\Pi)}$ holds for some universal constant $C > 0$ jin2023upper,zhan2022policy when $K \geq 3$. When $K = 2$, $d_{\mathrm{N}}(\Pi)$ is equal to the VC dimension, and $\kappa(\Pi) \leq 2.5\sqrt{d_{\mathrm{N}}(\Pi)}$ holds Haussler1995.

Furthermore, we make the following assumption for the $\epsilon$-Hamming covering number.

assumptionFor all $\epsilon \in (0, 1)$, $\mathbb{N}_H(\epsilon,\Pi) \le C\exp(D({\frac{1}{\epsilon}})^{\omega})$ for some constants $C,D > 0, 0 < \omega < 0.5$.

Then, we obtain the following upper bound for the expected simple regret of the PLAS strategy.

theorem[Upper bound] Suppose that Assumption (ref) holds. Then, for any $\zeta$ and any $P\in\mathcal{P}_\zeta$, the expected simple regret of the PLAS strategy satisfies \begin{align} &\sqrt{T}\mathbb{E}_P\Big[R\left(P\right)\left(\widehat{\pi}^{\mathrm{PLAS}}_T\right)\Big] \\ &\leq \begin{cases} \Big(108.8 \kappa(\Pi) + 870.4\Big)\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{\sum^K_{a=1}\left(\sigma^a(X)\right)^2}\right] + o(1) & \mathrm{if}\ K \geq 3\\ \Big(108.8 \kappa(\Pi) + 870.4\Big)\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{\left(\sigma^1(X) + \sigma^2(X)\right)^2}\right] + o(1) & \mathrm{if}\ K = 2 \end{cases}\quad \mathrm{as}\ \ T\to \infty, \end{align} where $\kappa(\Pi) = \begin{cases} 2.5\sqrt{d_{\mathrm{N}}(\Pi)} & \mathrm{if}\ K = 2\\ C\sqrt{\log(d)d_{\mathrm{N}}(\Pi)} & \mathrm{if}\ K \geq 3 \end{cases}$ holds for a universal constant $C > 0$.

From Lemma (ref) and the relationship between the entropy integral and the Natarajan dimension, the following theorem holds.

When $K = 2$, for $M = d_{\mathrm{VC}}(\Pi)$, the regret upper bound is given as $\Big(272 \sqrt{M} + 870.4\Big)\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{\sum^K_{a=1}\left(\sigma^a(X)\right)^2}\right] + o(1)$ as $T \to \infty$. When $K \geq 3$, for $M = d_{\mathrm{N}}(\Pi)$, the regret upper bound is given as $\Big(108.8 C\sqrt{\log(d)M} + 870.4\Big)\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{\sum^K_{a=1}\left(\sigma^a(X)\right)^2}\right] + o(1)$ as $T \to \infty$. Here, we note that the leading factor in these upper bounds are $\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{\log(d)M\sum^K_{a=1}\left(\sigma^a(X)\right)^2}\right]$ and $\mathop{\mathbb{E}}_{X \sim \zeta}\left[\sqrt{M\left(\sigma^a(X) + \sigma^2(X)\right)^2}\right]$, a product of the policy complexity ( Natarajan dimension $d_{\mathrm{N}}(\Pi)$) and outcome variances. This theorem implies that the leading factors align with the lower bounds with high probability.

figure[figure omitted — 233 chars of source]

Simulation Study

We conduct simulation studies to investigate the empirical performance of our proposed PLAS strategy. We compare the PLAS strategy with a combination of uniform sampling and policy learning, denoted as Uniform. The Uniform strategy assigns treatment arms with an equal ratio of $1/K$ and then applies policy learning. As a baseline method, we use the PLAS strategy with known variances, referred to as Oracle.

We consider a simple scenario with $K = 4$. We examine three cases for $d$ and $T$: $(d, T) = (2, 10000)$, $(d, T) = (5, 10000)$, and $(d, T) = (5, 25000)$. Let $X_{i}$ be the $i$-th dimension of $X$, and let $m_{i}$ and $v_{i}$ be its mean and variance, respectively. The mean $m_{i}$ is drawn from a uniform distribution with support $[-1, 1]$, and the variance $v_{i}$ is fixed at $1$. If $X_{(1)} > 0.5$ and $X_{(2)} > 0.5$, then $\mu^1(P)(X) = 5.00$ and $\mu^2(P)(X) = \mu^3(P)(X) = \mu^4(P)(X) = 4.50$; if $X_{(1)} < 0.5$ and $X_{(2)} > 0.5$, then $\mu^2(P)(X) = 5.00$ and $\mu^1(P)(X) = \mu^3(P)(X) = \mu^4(P)(X) = 4.50$; if $X_{(1)} > 0.5$ and $X_{(2)} < 0.5$, then $\mu^2(P)(X) = 5.00$ and $\mu^1(P)(X) = \mu^3(P)(X) = \mu^4(P)(X) = 4.50$; if $X_{(1)} < 0.5$ and $X_{(2)} < 0.5$, then $\mu^4(P)(X) = 5.00$ and $\mu^1(P)(X) = \mu^2(P)(X) = \mu^3(P)(X) = 4.50$.

We conduct $50$ independent trials to evaluate the performance of the strategies. The results are presented in Figure (ref) with three different settings for $d$ and $T$. The $y$-axis represents the policy value of the learned policy. From the experimental results, we confirm that our proposed strategy effectively improves the policy value.

Conclusion

In this study, we presented an adaptive experiment with policy learning. Our main contributions include the derivation of lower bounds for strategies, the development of the PLAS strategy, and the establishment of its regret upper bound. First, by utilizing the lower bounds developed by Kaufman2016complexity, we derived lower bounds for the expected simple regret, which depend on the variances of outcomes. Then, based on these lower bounds, we developed the PLAS strategy, which trains a policy at the end of the experiment. Lastly, we provided upper bounds for the regret of the PLAS strategy.

From a technical perspective, we demonstrated how to use Rademacher complexity for i.i.d. samples in an adaptive experiment with non-i.i.d. samples. We did not employ complexity measures for non-i.i.d. samples as presented by Rakhlin2015 and Foster2023. Instead, our technique relies on an approach used by Hahn2011, which we extended by incorporating sample splitting, also known as double machine learning Laan2008TheCA,ZhengWenjing2011CTME,ChernozhukovVictor2018Dmlf,hadad2019,Kato2020adaptive,Kato2021adr.

We also contributed to the literature on policy learning by providing a variance-dependent lower bound, which applies to observational studies with i.i.d. samples, and by discussing matching upper bounds. Our derived lower bound is distinct from that in AtheySusan2017EPL and more tightly depends on variances, necessitating the refinement of existing upper bounds.

Our next step is to tighten both the lower and upper bounds. When $K=2$, we showed that assigning each treatment arm a proportion based on standard deviations is optimal, consistent with existing works such as Neyman1934OnTT, glynn2004large, and Kaufman2016complexity. However, we found that assigning each treatment arm a proportion based on variances is optimal when $K \geq 3$. Other studies on fixed-budget BAI without contextual information, such as glynn2004large, Kaufman2016complexity, and kato2023locally, indicate that strategies with different sampling rules are optimal. Bridging the gap between our study and these existing studies remains an open issue.