Double machine learning and design in batch adaptive experiments

Double machine learning and design in batch adaptive experiments

Introduction

In sequential experimentation, we can use earlier observations to adjust our treatment allocation policy for subsequent observations and thereby gain improved estimation of causal effects in the overall study. For instance, for an experiment with one treatment arm and one control arm, neyman1934 showed that choosing the number of subjects in each arm to be proportional to the outcome standard deviation of that arm minimizes the variance of the treatment effect estimate based on the difference in means. While these standard deviations are unknown, they can be estimated using the initial data. Then the Neyman allocation can be approximated to improve the sample efficiency of the remainder of the experiment hahn2011adaptive,blackwell2022batch, zhao2023adaptive,dai2023clip.

We study a version of this design problem for an experiment divided into a small number of stages or batches. The design, or treatment assignment mechanism, can be updated adaptively for later batches based on the observations from earlier batches to improve the precision of the causal estimate computed at the end of the experiment. A salient feature of our setting is knowledge of pre-treatment covariates that can further improve precision. Thus, we conceptualize our design problem as choosing a propensity score for each batch. The propensity score specifies the probability that a subject receives treatment given their covariates (throughout, we consider the setting of a binary treatment). The propensity score is well known to be a key mathematical object to be estimated in the causal analysis of observational data (e.g. rosenbaum1983central). In a randomized experiment it is known and under the control of the investigator. Hence, it can be exploited for design.

Our specific design objective is to minimize an appropriate scalarization of the asymptotic covariance matrix of an estimator that efficiently pools information across all batches of the experiment. After describing our mathematical notation and setup in Section (ref), we present and study an oracle version of this pooled estimator in Section (ref). That oracle is given knowledge of some nuisance parameters, typically infinite-dimensional mean or variance functions. It pertains to a so-called “non-adaptive batch experiment" where treatment is assigned with possibly varying but nonrandom propensity scores across batches. In certain cases, our oracle pooled estimator asymptotically dominates the best possible alternative that aggregates single batch estimates, regardless of the per-batch propensities. This justifies designing for the pooled estimator instead of an aggregation-based alternative.

For such a design procedure to be useful, however, we must show that the targeted asymptotic precision of the oracle pooled estimator is in fact attainable in a batched experiment where the propensity scores used are adaptive (data-dependent) and the nuisance parameters need to be estimated. We address these challenges in Section (ref) by extending the framework of double machine learning formalized by chernozhukov2018double, hereafter DML, to what we call a “convergent split batch adaptive experiment," or CSBAE. In a CSBAE, observations in each batch are split into $K$ folds, and treatment is assigned in each fold according to an adaptive propensity score that only depends on observations from previous batches within the same fold. Within a given batch, the $K$ adaptive propensity scores are further required to converge to a common limit at rate $O_p(N^{-1/4})$ in root mean square (RMS). Our DML extension then shows that by plugging in estimated nuisance functions, we can construct a feasible estimator in a CSBAE that is asymptotically equivalent to the oracle pooled estimator computed on the limiting non-adaptive batch experiment. The nuisance function estimates only need to converge at the rate $o_p(N^{-1/4})$.

Section (ref) details a finite dimensional concave maximization procedure (Algorithm (ref)) that provably constructs a CSBAE for which the limiting propensities in each batch are sequentially optimal within a function class satisfying standard complexity conditions. Hence, we can effectively design for our pooled estimator in a batch adaptive experiment with theoretical guarantees that our final estimator will indeed attain the targeted optimal asymptotic precision, even with a fairly flexible propensity score learning method and nonparametric machine learning estimates for nuisance functions.

To the best of our knowledge, existing work either designs for a less efficient alternative to our pooled estimator, or discretizes the covariate space at the design and estimation stages to construct a feasible variant of the pooled estimator on a batch adaptive experiment. Our simulations in Section (ref) suggest that the latter approach in particular can lead to substantial precision losses that swamp any gains from design, even with a moderate number of continuous covariates. Thus, for the practitioner, we provide an end-to-end design and estimation procedure to efficiently handle continuous covariates in a batched experiment.

Related work

There has been substantial research interest in adaptive experiment designs in recent years. In many applications, treatment assignments are updated in an attempt to maximize the (expected) response values of either those in the experiment, as in adaptive bandit algorithms russo2018tutorial,hao2020adaptive, or those in the superpopulation from which the experimental subjects are assumed to arrive xu2018fully,kasy2021adaptive. Inference on data collected from these algorithms can be challenging since the treatment assignment rules often do not converge hadad2021confidence,zhang2020inference,zhang2021statistical. By contrast, in our setting where the goal is purely statistical (maximizing asymptotic precision of the treatment estimate), the design objective is a static propensity score to be learned consistently. An interesting direction for further study would be to design for a mixture of both statistical and non-statistical objectives. For example, one might expand the literature on tie-breaker designs owen2020optimizing,morrison2022optimality,li2023general,kluger2023kernel) to the setting of batched experiments.

The present work can be viewed as an extension of hahn2011adaptive in several directions. Those authors considered a two batch experiment to estimate the average treatment effect (ATE) as precisely as possible. Using data from the first batch to estimate variance functions, they estimate the asymptotic variance of a pooled version of the semiparametric efficient ATE estimator of hirano2003efficient for a coarsely discretized covariate. Then, they learn a propensity score for the second batch that approximately minimizes this variance. The covariate discretization ensures nuisance functions and optimal propensities can be estimated at parametric $O_p(N^{-1/2})$ rates without parametric assumptions. Consequently, a feasible version of the pooled estimator indeed attains the targeted asymptotic variance on the batch adaptive experiment. We generalize this pooling construction beyond the setting of ATE estimation and relax these rate requirements to those described in the previous section. This permits more efficient handling of continuous covariates through nonparametric nuisance function estimates and more flexible adaptive propensity scores.

Other approaches to extend the work of hahn2011adaptive include kato2020efficient, who consider an online setting where subjects from a stationary superpopulation enter one at a time, without batches. Similarly, the literature on covariate-adjusted response-adaptive (CARA) designs has focused on different but related objectives, both statistical and ethical zhang2007asymptotic,zhu2023covariate. In the batched setting, tabord-meehan2022stratification proposes a method to learn a variance-minimizing stratification of the covariate space of fixed size, avoiding the need to discretize the space prior to observing the data as in hahn2011adaptive. cytrynbaum2021designing showed that by performing a form of highly stratified treatment assignment called local randomization, consistent variance function estimates from the first batch make it possible to attain the semiparametric lower bound for ATE estimation in the second batch with optimal propensity score without having to estimate the conditional mean functions. Neither of these approaches, however, maintains the efficiency advantages of pooling. They also do not immediately extend beyond ATE estimation.

Setup and notation

Let $T \geqslant 2$ be the number of batches in the experiment. Each subject $i=1,\ldots,N_t$ in batch $t=1,\ldots,T$ has observed covariates $X_{ti}\in\mathbb{R}^d$ and potential outcomes $Y_{ti}(0),Y_{ti}(1)\in\mathbb{R}$. We place these in the vectors $S_{ti}=(X_{ti}^\top, Y_{ti}(0),Y_{ti}(1))^\top$ which are exogenous in our model. Let $Z_{ti} \in \{0,1\}$ be the binary treatment indicator for this subject, which is controlled by the investigator. Under the usual stable unit value treatment assumption (SUTVA), the observed outcome is

Then the available data for the subject is $W_{ti}=(X_{ti}^\top,Z_{ti},Y_{ti})^\top \in \mathcal{W}$. We assume that the vectors $S_{ti}\stackrel{\mathrm{iid}}{\sim} P^S$ for $t=1,\ldots,T$ and $i=1,\ldots,N_t$, for some distribution $P^S$. Appendix (ref) relaxes this assumption to permit certain forms of non-stationarity across batches, such as covariate shifts. It will be convenient to define the functions

for $z\in\{0,1\}$ and $x\in\mathcal{X}$. These expectations are taken under $P^S$.

Let $N = N_1 + \ldots +N_T$. Then as in hahn2011adaptive, che2023adaptive and others, we consider a proportional asymptotic regime

as $N \rightarrow \infty$. In settings where the batch sizes are fully controlled by the experimenter, it may be theoretically preferable to make initial batch sample sizes a vanishing fraction of the total sample size zhao2023adaptive. However, in many settings the batch sizes are exogenously constrained to satisfy (ref) unless observations are discarded.

For various $q \geqslant 1$ and probability measures $P$ on some space $\Omega$, it will be useful to consider function norms of the form \[ \|f\|_{q,P} = \left(\int |f(w)|^q \mathrm{d} P(w)\right)^{1/q} \] for $f\in L^q(P)$. We will use propensity scores denoted by $e(\cdot)$ with various subscripts. A propensity score $e(\cdot)$ specifies $e(x)=\textnormal{Pr}(Z=1 \mid X=x)$, the probability of treatment conditional on covariates. We will typically require propensity scores to lie in $\mathcal{F}_{\gamma}$, the set of all measurable functions on $\mathcal{X}$ taking on values in the interval $[\gamma,1-\gamma]$ for some $\gamma \in [0,1/2)$. We use $\|A\|$ to denote the square root of the sum of the squared entries of any vector, matrix, or tensor $A$. For any integer $p \geqslant 1$, $\mathbb{S}_+^p$ will denote the set of symmetric positive semidefinite $p \times p$ real matrices, and $\mathbb{S}_{++}^p$ will be the set of symmetric positive definite $p \times p$ real matrices. Finally, for any real vector $v$, we write $v^{\otimes 2}=vv^\top$.

We summarize the preceding requirements for the data generating process in Assumption (ref). Assumption (ref) does not impose any restrictions on the treatment assignment process, which will be discussed at length in subsequent sections.

Estimands and score equations

Consider the setting where $T=1$ (so we can drop the batch subscript $t$), and the observations $W_1,\ldots,W_N$ are i.i.d. Suppose additionally that (ref) holds along with the unconfoundedness assumption \[ (Y_i(0),Y_i(1)) \perp \!\!\! \perp Z_i \mid X_i, \quad i=1,\ldots,N. \] Then many popular causal estimands $\theta_0 \in \Theta \subseteq \mathbb{R}^p$ are identified by a score equation \[ \mathbb{E}[s(W;\theta_0,\nu_0,e_0)]=0. \] In this score equation, $\nu_0$ is a vector of possibly infinite-dimensional nuisance parameters lying in a nuisance set $\mathcal{N}$, and $e_0=e_0(\cdot):\mathcal{X} \rightarrow [0,1]$ is the propensity score. Following Section 3.1 of chernozhukov2018double, we will assume for simplicity that the score $s(\cdot)$ is linear in the sense that

for some $\gamma \in [0,1/2)$, $s_a(\cdot,\nu,e):\mathcal{W} \rightarrow \mathbb{R}^{p \times p}$, and $s_b(\cdot,\nu,e): \mathcal{W} \rightarrow \mathbb{R}^p$.

When $T>1$, propensity scores may vary across batches by design or external constraints. For any propensity $e=e(\cdot)$ and integrable function $f:\mathcal{W} \rightarrow \mathbb{R}$, we use the subscripted notation $\mathbb{E}_e[f(W)]=\int f(w) \mathrm{d} P_e(w)$ where $P_e=P_e^W$ is the distribution of $W=(X,Z,Y)=(X,Z,ZY(1)+(1-Z)Y(0))$ induced by $S=(X,Y(0),Y(1)) \sim P^S$ and $Z \!\mid\! X \sim \textnormal{Bern}(e(X))$ under the SUTVA assumption (ref). Further let $P^X$ be the marginal distribution of $X$ under $S \sim P^S$. Then we will require the following score equations to hold for some $\gamma \in [0,1/2)$ to identify the our causal estimand $\theta_0$:

Note that (ref) requires the score $s(\cdot)$ to have mean 0 when any propensity score $e'(\cdot) \in \mathcal{F}_{\gamma}$ is plugged in. This plug-in propensity $e'(\cdot)$ may differ from the propensity $e(\cdot) \in \mathcal{F}_{\gamma}$ used for treatment assignment in the experiment that generated the observations $W$. Such a robustness property is satisfied by definition so long as the score $s(\cdot)$ is doubly robust. It is required to ensure the validity of the pooled estimator that we propose in Section (ref). That estimator requires plugging in a mixture propensity score that averages propensities across all batches $t=1,\ldots,T$. While identification of $\theta_0$ within each batch is possible by only requiring (ref) to hold when $e(\cdot)=e'(\cdot)$, this will not be sufficient to ensure validity of our pooled estimator. We formally restate our requirements on identification of the estimand $\theta_0$ in Assumption (ref).

The first estimand we are motivated by is the ATE, given by $\theta_{0,\textnormal{ATE}}=\mathbb{E}[Y(1)-Y(0)]\in\mathbb{R}$ in our notation. For an investigator interested in modeling how the treatment effect varies with $X$, they may instead wish to estimate the regression parameter $\theta_{0,\textnormal{PL}} \in \mathbb{R}^p$ under a linear treatment effect assumption

See robinson1988root for background on semiparametric estimation of $\theta_{0,\textnormal{PL}}$ under (ref), which characterizes the well-known “partially linear model." We show next that both $\theta_{0,\textnormal{ATE}}$ and $\theta_{0,\textnormal{PL}}$ are identified by score functions that are linear in the sense of (ref) and robust in the sense of (ref), and hence identifiable according to Assumption (ref).

Some of our later results pertain to general estimands identified by Assumption (ref). Others will be specialized to the settings of Examples (ref) and (ref).

Oracle pooled and aggregate estimation (non-adaptive)

Here we propose and analyze an oracle estimator $\hat{\theta}^*$ of a generic estimand $\theta_0$ satisfying Assumption (ref) with score $s(\cdot)$. It is an oracle in the sense that it uses the unknown true value of the nuisance parameter $\nu_0$ from the score equations (ref). We discuss feasible estimation of $\theta_0$, including estimation of $\nu_0$, in Section (ref). The estimator $\hat{\theta}^*$ pools observations across all batches $t=1,\ldots,T$. We then prove a central limit theorem (CLT) for it in the setting of a non-adaptive batch experiment where treatment in each batch $t=1,\ldots,T$ is assigned according to a fixed (non-random) propensity score $e_t(\cdot)$:

Next, we compare the pooled estimator $\hat{\theta}^*$ to an alternative oracle that makes an optimal linear aggregation of per-batch estimates. It also satisfies a CLT, but we show that our pooling strategy dominates aggregation in terms of efficiency in the setting of Examples (ref) and (ref). We remark that some authors (e.g. tabord-meehan2022stratification) refer to this aggregation approach as pooling, but we reserve that term for pooling data, not estimators.

Pooled oracle estimator

The main idea behind the construction of our oracle estimator $\hat{\theta}^*$ is as follows. After collecting the observations from all batches $t=1,\ldots,T$ in a non-adaptive batch experiment, we ignore the batch structure and pool together the observations across batches. Now consider a random draw $W=(X,Z,Y)$ from these pooled observations $\{W_{ti} \mid 1 \leqslant t \leqslant T, 1 \leqslant i \leqslant N_t\}$. In the notation of Section (ref), it is straightforward to show that the distribution of $W$ is $P_{e_{0,N}}=\sum_{t=1}^T (N_t/N)P_{e_t}$, where $e_{0,N}(\cdot)$ is the mixture propensity score

It will also be helpful to define the limiting mixture propensity score $e_0(\cdot)$ under the proportional asymptotics (ref):

When a particular set of nonrandom batch propensities $e_1(\cdot),\ldots,e_T(\cdot)$ is relevant, we omit an additional subscript by letting $\mathbb{E}_{0,N}[f(W)]=\mathbb{E}_{e_{0,N}}[f(W)]$ and $\mathbb{E}_0[f(W)]=\mathbb{E}_{e_0}[f(W)]$. Using this notation, the oracle pooled estimator $\hat{\theta}^*$ is derived by solving the sample analogue of the mixture score equations $\mathbb{E}_{0,N}[s(W;\theta_0,\nu_0,e_{0,N})]=0$ for $\theta$:

This $\hat\theta^*$ is only defined when the matrix inverse in (ref) exists. That will be the case with probability tending to 1 so long as $\mathbb{E}_0[s_a(W;\nu_0,e_0)]$ is invertible, which we will require in our theoretical results. One such result is a CLT for $\hat{\theta}^*$. The proofs of all our technical results are provided in Appendix (ref).

For the estimands $\theta_{0,\textnormal{ATE}}$ and $\theta_{0,\textnormal{PL}}$, following (ref) we use the scores $s_{\textnormal{AIPW}}(\cdot)$ and $s_{\textnormal{EPL}}(\cdot)$, respectively to derive the oracle estimates

We now specialize the generic oracle CLT of Proposition (ref) to these two estimators under some regularity conditions.

The aggregated oracle estimator

When the covariate space $\mathcal{X}$ is finite, the oracle pooled estimator $\hat{\theta}^*_{\textnormal{AIPW}}$ is equivalent to an oracle variant of the estimator proposed by hahn2011adaptive. The complexity of constructing a feasible pooled estimator for an adaptive experiment in more general settings has led other authors to instead consider single batch estimators that can lose considerable efficiency. For instance, cytrynbaum2021designing proposes simply discarding the first batch in a two-batch experiment when computing the final estimate, which is clearly inadmissible in our proportional asymptotic regime (ref). Section 3.2 of tabord-meehan2022stratification suggests instead taking a linear aggregation of estimates computed separately on each batch, as described above. We will now show that even the best linearly aggregated oracle estimator is asymptotically dominated by our pooled estimators $\hat{\theta}^*_{\textnormal{AIPW}}$ and $\hat{\theta}^*_{\textnormal{EPL}}$. While tabord-meehan2022stratification hypothesized in their Appendix C.2 that this may be true for ATE estimation, they do not pursue this further as are unable to construct a feasible pooled estimator attaining the targeted oracle variance when batch propensities are chosen adaptively using their stratification trees. By contrast, our design approach will allow us to construct such an estimator using our extension of double machine learning in Section (ref).

Fix a non-adaptive batch experiment with batch propensities $e_1(\cdot),\ldots,e_T(\cdot)$. For each batch $t=1,\ldots,T$, consider an (oracle) estimator $\hat{\theta}_{t,\textnormal{AIPW}}^*$ for $\theta_{0,\textnormal{ATE}}$ computed by solving the empirical analogue of the score equations $\mathbb{E}_{e_t}[s_{\textnormal{AIPW}}(W;\theta_{0,\textnormal{ATE}},\nu_{0,\textnormal{AIPW}},e_t)]=0$ that averages only those observations in batch $t$: \[ \hat{\theta}_{t,\textnormal{AIPW}}^* = -\biggl(\frac{1}{N_t}\sum_{i=1}^{N_t} s_{\textnormal{AIPW},a}(W_{ti};\nu_{0,\textnormal{AIPW}},e_t)\biggr)^{-1} \biggl(\frac{1}{N_t}\sum_{i=1}^{N_t} s_{\textnormal{AIPW},b}(W_{ti};\nu_{0,\textnormal{AIPW}},e_t)\biggr). \] By applying Proposition (ref) with a single batch, for each $t=1,\ldots,T$ we obtain the CLT

where

Now, as stated in slud2018combining, the asymptotically unbiased linear combination of $\hat{\theta}_{1,\textnormal{AIPW}}^*,\ldots,\hat{\theta}_{T,\textnormal{AIPW}}^*$ with the smallest asymptotic covariance matrix with respect to the semidefinite ordering is the inverse covariance weighted estimator \[ \hat{\theta}_{\textnormal{AIPW}}^{*,(\mathrm{LA})} = \biggl(\,\sum_{t=1}^T \kappa_t V_{t,\textnormal{AIPW}}^{-1}\biggr)^{-1} \sum_{t=1}^T \kappa_t V_{t,\textnormal{AIPW}}^{-1}\hat{\theta}_{t,\textnormal{AIPW}}^*. \] This optimal linearly aggregated estimator $\hat{\theta}_{\textnormal{AIPW}}^{*,(\mathrm{LA})}$ satisfies the CLT \[ \sqrt{N}\bigl(\hat{\theta}_{\textnormal{AIPW}}^{*,(\mathrm{LA})}-\theta_{0,\textnormal{ATE}}\bigr) \stackrel{\mathrm{d\,\,}}\rightarrow \mathcal{N}(0,V_{\textnormal{AIPW}}^{(\mathrm{LA})}), \quad V_{\textnormal{AIPW}}^{(\mathrm{LA})} = \biggl(\sum_{t=1}^T \kappa_t V_{t,\textnormal{AIPW}}^{-1}\biggr)^{-1}. \]

We can similarly define the linearly aggregated oracle estimator $\hat{\theta}_{\textnormal{EPL}}^{*,(\mathrm{LA})}$ for $\theta_{0,\textnormal{PL}}$ based on combining per-batch estimates from the score $s_{\textnormal{EPL}}(\cdot)$, which satisfies a CLT $\sqrt{N}(\hat{\theta}^{*,(\mathrm{LA})}_{\textnormal{EPL}}-\theta_{0,\textnormal{PL}}) \stackrel{\mathrm{d\,\,}}\rightarrow \mathcal{N}(0,V_{\textnormal{EPL}}^{(\mathrm{LA})})$. Here $V_{\textnormal{EPL}}^{(\mathrm{LA})}$ is given by replacing $s_{\textnormal{AIPW}}(\cdot)$ with $s_{\textnormal{EPL}}(\cdot)$, $\nu_{0,\textnormal{AIPW}}$ with $\nu_{0,\textnormal{EPL}}$, and $\theta_{0,\textnormal{ATE}}$ with $\theta_{0,\textnormal{PL}}$ in the definition of $V_{\textnormal{AIPW}}^{(\mathrm{LA})}$. Our main result is then that regardless of the batch propensities, $\hat{\theta}_{\textnormal{AIPW}}^{*,(\mathrm{LA})}$ and $\hat{\theta}_{\textnormal{EPL}}^{*,(\mathrm{LA})}$ are asymptotically dominated by our pooled estimators $\hat{\theta}_{0,\textnormal{AIPW}}$ and $\hat{\theta}_{0,\textnormal{EPL}}$, respectively. This motivates our work in Section (ref) that designs for these estimators.

Feasible pooled estimation in batch adaptive experiments

The oracle estimator $\hat\theta^*$ of (ref) depends on nuisance parameters $\nu_0$ that are unknown in practice. Additionally, recall that our CLT for $\hat\theta^*$ (Proposition (ref)) holds for a non-adaptive batch experiment. Our goal is to choose propensities adaptively in each batch to improve precision. Therefore, we would like to develop a feasible estimator $\hat\theta$ that attains the targeted asymptotic variance for experiments where treatment is assigned adaptively, even when the nuisance parameters $\nu_0$ must be estimated.

As mentioned above, our construction of such a feasible estimator $\hat{\theta}$ is based on extending the double machine learning (DML) framework of chernozhukov2018double. The main requirements for $\hat{\theta}$ to have the same asymptotic variance as the corresponding oracle are convergence rate guarantees for both nuisance parameter estimates and the adaptive propensities. Our DML extension ensures that these rate requirements can be made sub-parametric, enabling the use of somewhat flexible machine learning methods.

The typical DML setting assumes access to a single sample $W_1,\ldots,W_N$ of i.i.d.\ observations. An example of this setting is a non-adaptive batch experiment with $T=1$ and propensity $e_0(\cdot)$. Then a standard DML estimator is based on two ingredients: a Neyman orthogonal score and cross-fitting. Neyman orthogonality of the score $s(\cdot)$ at $(\nu_0,e_0)$ means a local insensitivity of the score equations to perturbations in $(\nu_0,e_0)$ in any direction: \[ \frac{\partial}{\partial \lambda} \mathbb{E}_{e_0}\bigl[s(W_i;\theta_0,\nu_0+\lambda(\nu-\nu_0),e_0+\lambda(e-e_0))\bigr]=0, \quad \forall (\nu,e) \in \mathcal{N} \times \mathcal{F}_{\gamma}. \] It is well known that the scores $s_{\textnormal{AIPW}}(\cdot)$ and $s_{\textnormal{EPL}}(\cdot)$ are Neyman orthogonal chernozhukov2018double. Given a Neyman orthogonal score $s(\cdot)$, DML proceeds by constructing an estimator $\hat{\theta}$ by cross-fitting. In cross-fitting, the indices $1,\ldots,N$ are partitioned into $K$ (roughly) equally sized folds $\mathcal{I}_1,\ldots,\mathcal{I}_K$. Then $\hat{\theta}$ is computed as the solution to the empirical score equations $N^{-1} \sum_{k=1}^K \sum_{i \in \mathcal{I}_k} s(W_i;\hat{\theta},\hat{\nu}^{(-k)},\hat{e}^{(-k)}) = 0$, where for each $k=1,\ldots,K$, $\hat{\nu}^{(-k)}$ and $\hat{e}^{(-k)}(\cdot)$ are estimates of $\nu_0$ and $e_0(\cdot)$, respectively. Each pair $(\hat{\nu}^{(-k)},\hat{e}^{(-k)}(\cdot))$ depends only on the observations $\{W_i \mid i \notin \mathcal{I}_k\}$ outside fold $k$. The sample splitting ensures that for each $k=1,\ldots,K$, the estimates $(\hat{\nu}^{(-k)},\hat{e}^{(-k)})$ are independent of the observations in fold $k$. By the arguments of chernozhukov2018double, such independence is key to guarantee that the feasible estimator $\hat{\theta}$ is equivalent (up to first order asymptotics) to the oracle $\hat{\theta}^*$ solving $N^{-1} \sum_{k=1}^K \sum_{i \in \mathcal{I}_k} s(W_i;\hat{\theta},\nu_0,e_0) = 0$ (cf. (ref)), even when the estimates $(\hat{\nu}^{(-k)},\hat{e}^{(-k)})$ converge at sub-parametric rates.

To maintain this independence in an adaptive batched experiment, we require sample splitting at the design stage, as illustrated in Figure (ref), along with convergence of the adaptive propensity scores. Our notion of a convergent split batch adaptive experiment (CSBAE) in Definition (ref) formalizes this. The main idea is to split the observations in every batch $t=1,\ldots,T$ into $K$ folds. Re-using the notation above from the standard DML setting with $T=1$, we let $\mathcal{I}_k$ denote the set of batch and observation indices $(t,i)$ assigned to fold $k=1,\ldots,K$. Then the adaptive propensity used to assign treatment to a subject in batches $t=2,\ldots,T$ is allowed to only depend on observations in previous batches from the same fold as this subject. To ensure that the adaptivity does not introduce any additional variability (up to first-order asymptotics) into the final estimator, a CSBAE requires these adaptive propensities to converge to nonrandom limits $e_1(\cdot),\ldots,e_T(\cdot)$ at RMS rate $O_p(N^{-1/4})$. While this convergence requirement may appear restrictive, in Section (ref) we show how it can be ensured by design by solving an appropriate finite-dimensional concave maximization procedure. Moreover, the limiting propensity scores from this procedure will be provably optimal, in a sense we make more precise in Section (ref).

Given a CSBAE, our feasible estimator is

As in the standard (single batch) DML setting, for each $k=1,\ldots,K$, $\hat{\nu}^{(-k)}$ and $\hat{e}^{(-k)}(\cdot)$ are estimates of the nuisance parameters $\nu_0$ and the mixture propensity $e_{0,N}(\cdot)$ defined in (ref), respectively, that depend only on the observations $\{W_{ti} \mid (t,i) \notin \mathcal{I}_k\}$ outside fold $k$. These observations are fully independent of the observations in fold $k$ (across all batches $t=1,\ldots,T$) by the construction of a CSBAE. As in the single batch case, given $o_p(N^{-1/4})$ convergence of the estimators $\hat{\nu}^{(-k)}$ to $\nu_0$, this independence along with (ref) enable a DML-style argument that $\hat{\theta}$ is asymptotically equivalent to the oracle $\hat{\theta}^*$ under (ref) computed on a counterfactual non-adaptive batch experiment with propensities $e_1(\cdot),\ldots,e_T(\cdot)$. This argument proceeds by coupling the treatment indicators $Z_{ti}$ in the CSBAE with counterfactual treatment indicators $\tilde{Z}_{ti}=\bm{1}(U_{ti} \leqslant e_t(X_{ti}))$.

In Assumption (ref), the conditions (ref) and (ref) along with the inequalities (ref) through (ref) are direct extensions of Assumptions 3.1 and 3.2 in chernozhukov2018double for ordinary DML ($T=1$). The equations (ref) and (ref) are additional requirements that enable the dependence across batches in a CSBAE to be sufficiently weak so that $\hat{\theta}$ computed on the CSBAE is asymptotically equivalent to a version of it computed on the limiting non-adaptive batch experiment.

We can show that Assumption (ref) is satisfied for estimating $\theta_{0,\textnormal{ATE}}$ and $\theta_{0,\textnormal{PL}}$ with $s_{\textnormal{AIPW}}(\cdot)$ and $s_{\textnormal{EPL}}(\cdot)$, respectively under simple rate conditions on nuisance parameter and propensity estimation rates that mirror those in Section 5 of chernozhukov2018double for single-batch DML estimators. Then we apply Theorem (ref) to construct feasible pooled estimators $\hat{\theta}_{\textnormal{AIPW}}$ and $\hat{\theta}_{\textnormal{EPL}}$ as special cases of equation (ref), which attain the oracle asymptotic variances $V_{0,\textnormal{AIPW}}$ and $V_{0,\textnormal{EPL}}$ defined in (ref) and (ref), respectively.

We now compare the rate requirements in Corollaries (ref) and (ref) with those needed to prove a feasible CLT for the linearly aggregated estimators discussed in Section (ref). Consider a batch adaptive experiment without sample splitting, so that we assign treatment using the propensities $\hat{e}_1(\cdot),\ldots,\hat{e}_T(\cdot)$ where each $\hat{e}_t(\cdot)$ is possibly random but can only depend on the observations in batches $1,\ldots,t-1$, and converges to some nonrandom $e_t(\cdot)$ at any rate; in particular, this rate may be slower than the $O_p(N^{-1/4})$ rate of (ref). Then the standard DML results in Section 5 of chernozhukov2018double show that we can construct feasible estimators $\hat{\theta}_{t,\textnormal{AIPW}}$ and $\hat{\theta}_{t,\textnormal{EPL}}$ that are asymptotically equivalent to the oracle single-batch estimators $\hat{\theta}_{t,\textnormal{AIPW}}^*$ and $\hat{\theta}_{t,\textnormal{EPL}}^*$ in Section (ref), so long as we plug in the true propensity score $\hat{e}_t(\cdot)$ used to assign treatment and use cross-fitting at the estimation stage, even if all rates in Corollaries (ref) and (ref) are weakened to $o_p(1)$. Unfortunately, we cannot extend this construction to our pooled estimator $\hat{\theta}$ when $T \geqslant 2$, since $\hat{\theta}$ plugs in estimates of the mixture propensity $e_{0,N}(\cdot)$, which does not correspond (in general) to a propensity actually used for treatment in any batch. Our numerical studies in Section (ref) suggest, however, that the weaker rate requirements for a feasible CLT for the linearly aggregated estimators (compared to our pooled estimators) make little difference in practice. Indeed, there we also see finite sample advantages to pooling, beyond those predicted by the asymptotics.

Batch adaptive learning of the optimal propensity score

We now discuss how to learn adaptive propensity scores $\hat{e}_t^{(k)}(\cdot)$ that satisfy (ref) with limiting propensity scores $e_t(\cdot)$ that maximize asymptotic precision of the final estimators $\hat{\theta}_{\textnormal{AIPW}}$ and $\hat{\theta}_{\textnormal{EPL}}$ constructed in the previous section, as measured by their asymptotic covariance matrices $V_{0,\textnormal{AIPW}}$ and $V_{0,\textnormal{EPL}}$. This generates a CSBAE on which, by Theorem (ref), feasible estimators $\hat{\theta}_{0,\textnormal{AIPW}}$ and $\hat{\theta}_{0,\textnormal{EPL}}$ achieve the targeted asymptotic variances $V_{0,\textnormal{AIPW}}$ and $V_{0,\textnormal{EPL}}$.

While $V_{0,\textnormal{AIPW}}$ is scalar and so there is no ambiguity in what it means to maximize asymptotic precision of $\hat{\theta}_{\textnormal{AIPW}}$, when $\theta_{0,\textnormal{PL}}$ is multivariate ($p>1$), $V_{0,\textnormal{EPL}}$ is a matrix. To handle the multivariate setting, we follow classical literature on experiment design in regression models by scalarizing an information matrix $\mathcal{I} \in \mathbb{S}_+^p$ (typically an inverse covariance matrix) using an information function $\Psi:\mathbb{S}_+^p \rightarrow \mathbb{R} \cup \{-\infty\}$. See textbooks such as atkinson2007optimum,pukelsheim2006optimal for more background. We generically write $\mathcal{I}=\mathcal{I}(e,\eta)$ to emphasize that in our setting, $\mathcal{I}$ will be indexed by a propensity score $e=e(\cdot)$ in some function class $\mathcal{F}_*$ along with some unknown nuisance functions $\eta=\eta(\cdot)$ to be estimated. Then the design objective is to learn an optimal propensity $e^*(\cdot)$, in the sense of maximizing $\Psi(\mathcal{I})$:

The nuisance functions $\eta$ in the information matrix are distinct from the nuisance functions $\nu$ in the score function. Examples of $\eta$ for information matrices based on $V_{0,\textnormal{AIPW}}$ and $V_{0,\textnormal{EPL}}$ are given in Section (ref).

Generic convergence rates with concave maximization

We start by considering a feasible procedure to learn the propensity $e^*(\cdot)$ of (ref) when the information matrix $\mathcal{I}$ and the function class $\mathcal{F}_*$ being optimized over generically satisfy Assumption (ref) below. An explanation of how to apply this generic procedure to construct an appropriate CSBAE that designs for the estimators $\hat{\theta}_{\textnormal{AIPW}}$ and $\hat{\theta}_{\textnormal{EPL}}$ is deferred to Section (ref). Here we instead focus on a precise exposition of our propensity learning procedure and the technical assumptions needed to ensure convergence, without explicitly invoking any notation or setup from previous sections. We begin with some generic structure on the information matrix $\mathcal{I}$ and the function class $\mathcal{F}_*$. Roughly, we require the information matrix to be strongly concave in $e(\cdot)$ and the function class $\mathcal{F}_*$ to be not too complex. We also allow for interval constraints on the per-batch expected proportion of subjects treated.

Similar to the idea of empirical risk minimization in supervised learning (ERM, e.g. vapnik1991principles and montanari2022universality), when Assumption (ref) is satisfied, our learning procedure replaces the unknown population expectation appearing in the objective (ref) by a sample average over observations $X_1,\ldots,X_n \stackrel{\mathrm{iid}}{\sim} P$, where the generic sample size $n$ diverges. In our designed CSBAE, these observations will be the covariates of those subjects in a given batch $t=1,\ldots,T$ of a CSBAE within a given fold $k=1,\ldots,K$, so $n$ will be identified with the quantity $n_{t,k}$ of Definition (ref). The main computational procedure is then a finite dimensional optimization over the values of the propensity score at the $n$ points $X_1,\ldots,X_n$:

Here $\hat{\eta}(\cdot)$ is an estimate of the nuisance function $\eta(\cdot)$ in (ref), and the optimization set $F_n=F_n(m_L,m_H) \subseteq \mathbb{R}^n$ is defined by

where $(m_L,m_H)$ are budget constraints as in Assumption (ref) and $E_n$ is as in the numbered condition (ref) of that assumption.

We convert the vector $(\hat{e}_1,\ldots,\hat{e}_n)$ in (ref) to a propensity score $\hat{e}(\cdot)$ by taking $\hat{e}(\cdot)$ to be any member of the base propensity class $\mathcal{E}$ of Assumption (ref) with $\hat{e}(X_i)=\hat{e}_i$ for each $i=1,\ldots,n$. The existence of such a propensity $\hat{e}(\cdot)$ is guaranteed by the numbered condition (ref) in Assumption (ref). Then as in the ERM literature, we use empirical process arguments to guarantee the learned propensity $\hat{e}(\cdot)$ converges to $e^*(\cdot)$ at rate $O_p(n^{-1/4})$ under appropriate restrictions on the complexity of the base propensity class $\mathcal{E}$. Our main such restriction is the finite entropy integral requirement in (ref). Examples of function classes satisfying this condition can be found in the literature on empirical processes. A sufficient condition is that $\log \mathcal{N}(\epsilon,\mathcal{E},L^2(P_n)) \leqslant K\epsilon^{-2+\delta}$ for some $\delta >0$, some $K<\infty$, and all $\epsilon > 0$. Examples (ref) through (ref) below show that this requirement is loose enough to admit some relatively rich function classes.

Next, we construct sets $E_n$ that satisfy condition (ref) of Assumption (ref) for the base propensity classes in Examples (ref) to (ref). For the set of monotone functions in one dimension (Example (ref)) we can take \[ E_n = \{(e_1,\ldots,e_n) \mid 0 \leqslant e_{\pi(1)} \leqslant e_{\pi(2)} \leqslant \cdots \leqslant e_{\pi(n)} \leqslant 1\} \] where $\pi(\cdot)$ is the inverse of the function that maps each $i \in \{1,\ldots,n\}$ to the rank of $x_i$ among $x_1,\ldots,x_n$ (with any ties broken in some deterministic way). For the set of $L$-Lipschitz functions (Example (ref)) we can take \[ E_n = \bigl\{(e_1,\ldots,e_n) \bigm| |e_{\pi(i)}-e_{\pi(i-1)}| \leqslant L(x_{\pi(i)}-x_{\pi(i-1)}),\, i=2,\dots,n\bigr\}. \] For the parametric class in Example (ref) we can take \[ E_n = \bigl\{(\theta^{\top}\xi(x_1),\ldots,\theta^{\top}\xi(x_n)) \mid \theta^{\top}\bm{1}_p \leqslant 1, \theta \succcurlyeq 0\bigr\}. \] Finally, for the class (ref) in Example (ref) we take \[ E_n = \biggl\{\,\sum_{i=1}^m \omega_i\bigl(\textnormal{expit}(\theta_i^{\top}\varphi(x_1)),\ldots,\textnormal{expit}(\theta_i^{\top}\varphi(x_n))\bigr) \Bigm| \omega_i \geqslant 0, \sum_{i=1}^m \omega_i \leqslant 1 \biggr\}. \]

The convergence rates of $\hat{e}(\cdot)$ to $e^*(\cdot)$ will be proven using strong concavity of the design objective (ref) on the space $\mathcal{F}^*$. This is ensured by Assumption (ref) along with the following conditions on the information function $\Psi(\cdot)$:

We can show that Assumption (ref) is satisfied by two common information functions: the “$A$-optimality" function $\Psi_a(\cdot) = -\textnormal{tr}((\cdot)^{-1})$ with $\Psi_a(M) :=-\infty$ whenever $M$ is singular, and the “$D$-optimality" function $\Psi_d(\cdot) = \log(\det(\cdot))$. The $A$-optimality criterion corresponds to minimizing the average (asymptotic) variance of the components of the estimand, while $D$-optimality corresponds to minimizing the volume of the ellipsoid spanned by the columns of the (asymptotic) covariance matrix.

There are some common information functions that do not satisfy Assumption (ref). For example, the “$E$-optimality" function $\Psi_e(\cdot) = \lambda_{\min}(\cdot)$, where $\lambda_{\min}(M)$ refers to the smallest eigenvalue of $M \in \mathbb{R}^{p \times p}$, is not differentiable. Similarly the function $\Psi_c(\cdot) = -c^{\top}(\cdot)^{-1}c$ (for some fixed $c \in \mathbb{R}^p$), corresponding to “$c$-optimality," does not satisfy condition (c). We leave open the question of whether the $O_p(n^{-1/4})$ convergence rate of $\hat{e}(\cdot)$ to $e^*(\cdot)$ in Lemma (ref) can be extended to these (and other) information functions using different techniques, and now prove this rate under Assumptions (ref) and (ref).

Convergence of batch adaptive designs

We now leverage Lemma (ref) to develop a procedure (Algorithm (ref)) that can learn adaptive propensities $\hat{e}_t^{(k)}(\cdot)$ with the convergence guarantees (ref) so that when used for treatment assignment, they lead to a CSBAE with limiting propensities that optimize objectives of the form (ref) with information matrices based on $V_{0,\textnormal{AIPW}}^{-1}$ and $V_{0,\textnormal{EPL}}^{-1}$. This shows we can effectively design for the estimators $\hat{\theta}_{\textnormal{AIPW}}$ and $\hat{\theta}_{\textnormal{EPL}}$.

For simplicity, we assume treatment in the first batch is assigned according to a non-random propensity $e_1(\cdot) \in \mathcal{F}_{\epsilon_1}$ for some $\epsilon_1>0$. We let $\hat{e}_1^{(k)}(\cdot)=e_1(\cdot), k=1,\ldots,K$. Then for later batches $t=2,\ldots,T$, the target propensities are taken to be one of the following, for an information function $\Psi(\cdot)$ satisfying Assumption (ref):

Above, $V_{0:t,\textnormal{AIPW}}$ and $V_{0:t,\textnormal{EPL}}$ are the asymptotic variances of the oracle pooled estimators $\hat{\theta}_{\textnormal{AIPW}}$ and $\hat{\theta}_{\textnormal{EPL}}$ of (ref) and (ref), respectively, when computed using observations in a non-adaptive batch experiment with only $t$ batches and propensities $e_1(\cdot),\ldots,e_t(\cdot)$. By (ref) we can compute

where $\eta_{0,\textnormal{AIPW}}(x)$ includes the components $(v_0(0,x),v_0(1,x),\tau_0(x),\theta_0)$. Similarly, by (ref) we have

where $\eta_{0,\textnormal{EPL}}(x)$ includes the components $(v_0(0,x),v_0(1,x))$. In both of the preceding equations, the dependence on the batch $t$ propensity score $e_t(\cdot)$ is through the mixture $e_{0:t}(\cdot)$ given by

Finally, the optimization set $\mathcal{F}_{*,t}$ in (ref) is

which satisfies all the conditions in Assumption (ref) with covariate distribution $P^X$, budget constraints $(m_{L,t},m_{H,t})$, and base propensity class $\mathcal{E}$.

We do not target the final covariances $V_{0,\textnormal{AIPW}}=V_{0:T,\textnormal{AIPW}}$ and $V_{0,\textnormal{EPL}}=V_{0:T,\textnormal{EPL}}$ in our discussion here since the sample sizes and budget constraints in future batches may not be known. If they are known in advance, then we can learn propensities for all future batches simultaneously at the time batch 2 covariates are observed, and indeed target $V_{0:T,\textnormal{AIPW}}$ or $V_{0:T,\textnormal{EPL}}$ at that stage.

Now suppose we split our observations into $K$ folds as in a CSBAE, and for notational simplicity we re-index the covariates in each batch $t=1,\ldots,T$, fold $k=1,\ldots,K$ as $X_{t1}^{(k)},\ldots,X_{tn_{t,k}}^{(k)}$. Then following (ref), we can estimate $e_{t,\textnormal{AIPW}}^*(\cdot)$ for each batch $t \geqslant 2$ within each fold $k=1,\ldots,K$ by computing

Here $F_{n_{t,k}}=F_{n_{t,k}}(m_{L,t},m_{H,t})$ is defined as in (ref), and the estimate $\hat{V}_{0:t,\textnormal{AIPW}}^{(k)}$ of $V_{0:t,\textnormal{AIPW}}$ is given by

where each estimate $\hat{v}_{1:(t-1)}^{(k)}(z,\cdot)$ of the variance function $v_0(z,\cdot)$ is computed using only the observations from batches $u=1,\ldots,t-1$ within fold $k$. Note that any plug-in estimate of $\tau_0(\cdot)$ and $\theta_0$ does not affect the optimization (ref) and so can be omitted. The dependence of $\hat{V}_{0:t,\textnormal{AIPW}}^{(k)}$ on the optimization variables $e_1,\ldots,e_{n_{t,k}}$ is through the mixture quantities

where for $u=1,\ldots,t-1$, $\hat{e}_u^{(k)}(\cdot)$ is the (possibly adaptive) propensity used to assign treatment in batch $u$, fold $k$. By comparing (ref) and (ref), we see that for each batch $u=1,\ldots,t$, $N_u/N$ is being used as a plug-in estimate of $\kappa_u$ and the adaptive propensity score $\hat{e}_u^{(k)}(\cdot)$ is used as a plug-in estimate of its limit $e_u(\cdot)$. Finally, as in the conclusion of Lemma (ref), the learned adaptive propensity $\hat{e}_{t,\textnormal{AIPW}}^{(k)}(\cdot)$ to be used for treatment assignment in batch $t$, fold $k$ is taken to be any choice in the predetermined base collection $\mathcal{E}$ with $\hat{e}_{t,\textnormal{AIPW}}^{(k)}\left(X_{ti}^{(k)}\right)=\hat{e}_{ti,\textnormal{AIPW}}^{(k)}$ for all $i=1,\ldots,n_{t,k}$.

Learning $e_{t,\textnormal{EPL}}^*(\cdot)$ is exactly analogous; first we compute

where

Then we assign treatment with any propensity $\hat{e}_{t,\textnormal{EPL}}^{(k)}(\cdot)$ in the base propensity class $\mathcal{E}$ satisfying $\hat{e}_{t,\textnormal{EPL}}\big(X_{ti}^{(k)}\big)=\hat{e}_{ti,\textnormal{EPL}}^{(k)}$ for all $i=1,\ldots,n_{t,k}$.

The main additional regularity condition required to ensure the adaptive propensities $\hat{e}_{t,\textnormal{AIPW}}^{(k)}(\cdot)$ and $\hat{e}_{t,\textnormal{EPL}}^{(k)}$ above converge at the desired $O_p(N^{-1/4})$ RMS rate to $e_{t,\textnormal{AIPW}}^*$ and $e_{t,\textnormal{EPL}}^*$ of (ref) is the same rate of convergence in the estimates $\hat{\eta}^{(k)}_{\textnormal{AIPW}}(\cdot)$ and $\hat{\eta}^{(k)}_{\textnormal{EPL}}(\cdot)$ of the nuisance parameters $\eta_{0,\textnormal{AIPW}}(\cdot)$ and $\eta_{0,\textnormal{EPL}}(\cdot)$. We also strengthen the sample size asymptotics (ref) by requiring

Numerical simulations

We implement Algorithm (ref) to construct some synthetic CSBAE's that illustrate the finite sample performance of our proposed methods. For simplicity we consider $T=2$ batches throughout. Our evaluation metric is the average mean squared error (AMSE) of the estimators $\hat{\theta}_{\textnormal{AIPW}}$ and $\hat{\theta}_{\textnormal{EPL}}$ computed at the end of each CSBAE. As a baseline, we also compute feasible variants of the linearly aggregated estimators $\hat{\theta}_{\textnormal{AIPW}}^{(\mathrm{LA})}$ and $\hat{\theta}_{\textnormal{EPL}}^{(\mathrm{LA})}$ computed on a “simple RCT": that is, a non-adaptive batch experiment with a constant propensity score in each batch. We additionally consider the approach of hahn2011adaptive for design and estimation of $\theta_{0,\textnormal{ATE}}$. This is equivalent to using Algorithm (ref) for design (without sample splitting, i.e.\ $K=1$) and using the pooled $\hat{\theta}_{\textnormal{AIPW}}$ as the final estimator, but with the covariates $X$ replaced everywhere by a coarse discretization $S=S(X)$. As a hybrid we also consider using the discretized covariates $S$ for design but computing the final estimates $\hat{\theta}_{\textnormal{AIPW}}$ and $\hat{\theta}_{\textnormal{EPL}}$ using the full original covariate $X$. To separately attribute efficiency gains to design and pooling, we also consider the linearly aggregated estimators $\hat{\theta}_{\textnormal{AIPW}}^{(\mathrm{LA})}$ and $\hat{\theta}_{\textnormal{EPL}}^{(\mathrm{LA})}$ when a modification of Algorithm (ref) that targets $V_{2,\textnormal{AIPW}}$ and $V_{2,\textnormal{EPL}}$ (the asymptotic variances of the estimators $\hat{\theta}_{2,\textnormal{AIPW}}$ and $\hat{\theta}_{2,\textnormal{EPL}}$ given in Section (ref), depending only on observations in batch 2) is used for design.

We consider four data generating processes (DGPs), distinguished by whether the covariate dimension $d$ is 1 or 10 and whether the conditional variance functions $v_0(\cdot,\cdot)$ are homoskedastic (with $v_0(0,x)=v_0(1,x)=1$ for all $x \in \mathcal{X}$) or heteroskedastic (with $v_0(0,x)=v_0(1,x)/2=\exp((\bm{1}_d^{\top}x)/(2\sqrt{d}))$. The scaling by the covariate dimension $d$ in the heteroskedastic variance functions ensures the variance in $v_0(z,X)$ is independent of the covariate dimension $d$. In all of the DGPs the covariates are i.i.d. spherical Gaussian, i.e.\ $P^X=\mathcal{N}(0,I_d)$. The outcome mean functions are taken to be $m_0(0,x)=m_0(1,x)=\bm{1}_d^{\top}x$ for $\bm{1}_d=(1,\ldots,1)' \in \mathbb{R}^d$. For estimating $\theta_{0,\textnormal{PL}}$ we use the basis functions $\psi(x)=(1,x^{\top})^{\top} \in \mathbb{R}^p$ where $p=d+1$. Note $\theta_{0,\textnormal{ATE}}=\theta_{0,\textnormal{PL}}=0$. The potential outcomes $Y(0),Y(1)$ are generated as follows:

For each DGP we run Algorithm (ref) with $K=2$ folds, batch sample sizes $N_1=N_2=1000$, information function $\Psi=\Psi_a(\cdot)$ (corresponding to $A$-optimality), and treatment fraction constraints $m_{L,t}=m_{H,t}=0.2$ for $t=1,2$. In Appendix (ref), we present additional simulation results where $m_{L,1}=m_{H,1}$ remain at 0.2 but $m_{L,2}=m_{H,2}=0.4$, so that the treatment budget for the second batch has increased. The initial propensity score $e_1(\cdot)$ is taken to be constant (i.e. $e_1(x)=0.2$ for all $x \in \mathcal{X}$), and the base propensity class $\mathcal{E}$ is the set of all $1$-Lipschitz functions taking values in $[0,1]$ when $d=1$ (cf. Example (ref)). When $d=10$, we take $\mathcal{E}$ as in (ref), with $\{\theta_1,\ldots,\theta_m\}$ the collection of vectors $(a_1,\ldots,a_{11})' \in \mathbb{R}^{11}$ with each coordinate $a_i \in \{-2,-1,0,1,2\}$ and no more than two of the $a_i$'s nonzero.

The discretization used to implement the approach of hahn2011adaptive partitions $\mathbb{R}^d$ into four bins based on the quartiles of $\bm{1}_d^{\top}X$. Note that this partition is along the single dimension along which the variance functions $v_0(\cdot,\cdot)$ vary in the heteroskedastic DGPs, so we would expect this to perform better than in practice, where the structure of the variance functions is not known (it could possibly be learned, as in tabord-meehan2022stratification). The choice of four bins is based on the experiments of hahn2011adaptive, which find minimal performance difference between two and six bins for the DGP's they consider. We denote their “binned" AIPW estimator by $\hat{\theta}_{\textnormal{AIPW}}^{(\text{bin})}$. Recall, as indicated above, that this is equivalent to $\hat{\theta}_{\textnormal{AIPW}}$ when the only available covariate is the binned $S(X)$ and no cross-fitting is used. For $\hat{\theta}_{\textnormal{AIPW}}^{(\text{bin})}$ the conditional means $\tilde{m}_0(z,s) = \mathbb{E}(Y(z) \mid S=s), z=0,1$ are estimated nonparametrically by sample outcome means among all units with $Z=z$ and $S=s$ in the appropriate batch(es) and fold(s). Similarly, the conditional variances $\tilde{v}_0(z,s)=\textnormal{Var}(Y(z) \mid S=s)$ are estimated by sample outcome variances.

For all simulations, the concave maximizations are performed using the CVXR software fu2017cvxr and the MOSEK solver mosek. All estimates $\hat{m}(z,\cdot)$ of the mean function $m_0(z,\cdot), z=0,1$ are computed by fitting a generalized additive model (GAM) to the outcomes $Y$ and covariates $X$ from the observations with treatment indicators equal to $z$ in the appropriate fold(s) and batch(es). The GAMs use a thin-plate regression spline basis wood2003thin-plate and the degrees of freedom are chosen using the generalized cross-validation procedure implemented in the mgcv package in R wood2004stable. Variance function estimates $\hat{v}(z,\cdot)$ are computed by first computing $\hat{m}(z,\cdot)$ as above on the appropriate observations, then fitting a GAM on these same observations to predict the squared residuals $(Y-\hat{m}(z,X))^2$ from $X$.

Average treatment effect

Table (ref) shows the performance of the various design and estimation procedures for $\theta_{0,\textnormal{ATE}}$ that we consider. Each entry is a “relative efficiency": that is, a ratio of the MSE of the baseline approach (which computes the linearly aggregated $\hat{\theta}_{\textnormal{AIPW}}^{(\mathrm{LA})}$ on a simple RCT) to the MSE of the relevant approach. The simulated relative efficiencies in the table estimate these MSE's by averaging the squared error of each estimator over 1,000 simulations. The 90% confidence intervals for the true finite sample relative efficiency are computed using 10,000 bootstrap replications of these 1,000 simulations. Finally, the asymptotic relative efficiencies in Table (ref) are computed by estimating the asymptotic variance of each estimator using the appropriate formula, i.e. $V_{0,\textnormal{AIPW}}$ for $\hat{\theta}_{\textnormal{AIPW}}$ and $V_{\textnormal{AIPW}}^{(\mathrm{LA})}$ for $\hat{\theta}_{\textnormal{AIPW}}^{(\mathrm{LA})}$. The computation is based on a non-adaptive batch experiment with second batch propensity $e_2(\cdot)$ equal to the average of the learned propensities from the relevant approach across the 1,000 simulations, as a closed form solution for the limiting $e_2(\cdot)$ is not easily obtained in general. All expectations over the covariate distribution are computed using Monte Carlo integration.

For both of the homoskedastic DGPs, it is straightforward to show using Jensen's inequality that $e_2^*(x)=0.2$ for all $x$. It can be further shown that the unpooled and pooled estimators are asymptotically equivalent. Thus, for the homoskedastic DGPs, there is no asymptotic efficiency gain to be had. With unequal budget constraints there is some asymptotic benefit to pooling with homoskedastic variance functions (Appendix (ref)); however the optimal design will still be the simple RCT. Nonetheless, we do observe some finite sample benefits to pooling in Table (ref). For example, the simulated relative efficiency of $\hat{\theta}_{\textnormal{AIPW}}$ on the simple RCT is significantly larger than 1 for both $d=1$ and $d=10$. We attribute this to improved nuisance estimates for the pooled estimator, as discussed further in Appendix (ref). As one might expect, this finite sample improvement from pooling is apparently offset by variance in both the flexible and binned design procedures. Still, in the homoskedastic DGPs, our adaptive approaches do not show significant finite sample performance decline relative to the baseline, which here is an oracle. With unequal treatment constraints, we further obtain asymptotic efficiency gains from pooling (Appendix (ref)).

We notice that using the discretized covariate $S(X)$ in place of $X$ for both design and estimation, as in hahn2011adaptive, leads to a substantial loss of efficiency. Indeed, when $d=10$ the (asymptotic and simulated) variance of the estimator $\hat{\theta}_{\textnormal{AIPW}}^{(\text{bin})}$ is more than double that of our baseline under the homoskedastic DGP, both asymptotically and in our finite sample simulations. This efficiency loss occurs because the discretized $S(X)$ explains much less of the variation in the potential outcomes $Y(z)$ than the original $X$. We expect greater precision losses from this discretization at the estimation stage when the variance functions $v_0(z,\cdot)$ vary substantially within the strata defined by $S(X)$.

For the heteroskedastic DGPs, we see modest asymptotic efficiency gains from both pooling and design. Design using the flexible base propensity class leads to about a 2.4% asymptotic efficiency gain for both $d=1$ and $d=10$, while pooling provides an additional 1--2% gain. These small asymptotic gains appear to be largely canceled out at our sample sizes by the finite sample variability in learning propensity scores, limiting the net finite sample gains from design. Of course, with greater heteroskedasticity and/or differences between $v_0(0,\cdot)$ and $v_0(1,\cdot)$, we would expect greater efficiency gains from design; in our simulations we have chosen to keep these differences within common ranges in social science studies as per blackwell2022batch.

Partially linear model

Unlike for estimating $\theta_{0,\textnormal{ATE}}$, for estimating $\theta_{0,\textnormal{PL}}$, we see clear efficiency gains over the baseline from design when $d=1$ (Table (ref)). For instance, the linearly aggregated estimator $\hat{\theta}_{\textnormal{EPL}}^{(\mathrm{LA})}$ exhibits a 5.6% asymptotic efficiency gain as a result of the flexible design; replacing this with the pooled estimator $\hat{\theta}_{\textnormal{EPL}}$ then yields a total asymptotic gain of 10.0% over the baseline, even in the homoskedastic DGP. The analogous gains for the heteroskedastic DGP are slightly larger. We once again observe a substantial finite sample benefit to pooling, with the simulated relative efficiency of the approaches using the pooled $\hat{\theta}_{\textnormal{EPL}}$ tending to be larger than the asymptotic relative efficiency. We attribute this to both improved use of nuisance estimates by the pooled estimator (as in the ATE case) as well as a more fundamental finite sample efficiency boost due to the fact that the asymptotic variance $V_{0,\textnormal{EPL}}$ is not exact for the oracle pooled $\hat{\theta}_{\textnormal{EPL}}^*$ in finite samples, whereas $V_{0,\textnormal{AIPW}}$ is exact for $\hat{\theta}_{\textnormal{AIPW}}^*$; see Appendix (ref). When $d=10$, design introduces some more salient finite sample variance from the errors in the concave maximization procedure. This is offset in both the homoskedastic and heteroskedastic DGPs by the asymptotic gains from the flexible design, so that when $\hat{\theta}_{\textnormal{EPL}}$ is used, the flexible design ultimately performs similarly to the simple RCT in finite samples.

Even if we use $\hat{\theta}_{\textnormal{EPL}}$ and hence the original covariate(s) $X$ in the final estimation step, we see that the “binned" design which uses only $S(X)$ for choosing the second batch propensity score struggles to learn a substantially better propensity score than the simple RCT in all DGPs. Indeed, in the heteroskedastic DGP with $d=1$, we see a 2.1% asymptotic efficiency loss relative to the baseline from using the binned design. By comparison, there is an 11.2% asymptotic efficiency gain from using the flexible design. The efficiency loss can occur with the binned design because the objective in (ref) changes when working in terms of the discretized covariate $S(X)$ instead of $X$. In other words, even though the simple RCT propensity $e_2(x)=0.2$ is within the class of propensities that can be chosen by the binned design, it is worse according to the binned objective based on $S(X)$, but not according to the objective based on the original $X$. \\

Discussion

We view our primary technical contribution in this paper to be a careful extension of the double machine learning framework that enables both estimation and design in batched experiments based on pooled treatment effect estimators. This allows the investigator to take advantage of the efficiency gains from pooling and design without needing to make strong parametric assumptions or to discretize their covariates. As our numerical study in Section (ref) shows, the latter can more than wipe out any efficiency gains from design.

Related to our work is the extensive literature on combining observational data with a (single batch) randomized experiment. In that setting a primary concern is mitigating bias from unobserved confounders in the observational data rosenman2021designing,gagnon2023precise. By contrast, in our setting unconfoundedness holds by design in each batch of the experiment. It would be useful to examine if the ideas from the present work can be extended to the observational setting where confounding bias is a concern. \\

Acknowledgments: The authors thank Stefan Wager, Lihua Lei, and Kevin Guo for comments that improved the content of this paper. {H.L. was partially supported by the Stanford Interdisciplinary Graduate Fellowship (SIGF). This work was also supported by the NSF under grant DMS-2152780.