SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ting Li $^{1}$  Chengchun Shi $^{2}$  Jianing Wang $^{1}$  Fan Zhou $^{1}$  Hongtu Zhu $^{3}$ *
$^{1}$ School of Statistics and Management, Shanghai University of Finance and Economics
$^{2}$ Department of Statistics, London School of Economics and Political Science
$^{3}$ Department of Biostatistics, University of North Carolina at Chapel Hill
tingli@mail.shufe.edu.cn, c.shi7@lse.ac.uk, jianing.wang@163.sufe.edu.cn
zhoufan@mail.shufe.edu.cn,htzhu@email.unc.edu
A/B testing is critical for modern technological companies to evaluate the effectiveness of newly developed products against standard baselines. This paper studies optimal designs that aim to maximize the amount of information obtained from online experiments to estimate treatment effects accurately. We propose three optimal allocation strategies in a dynamic setting where treatments are sequentially assigned over time. These strategies are designed to minimize the variance of the treatment effect estimator when data follow a non-Markov decision process or a (time-varying) Markov decision process. We further develop estimation procedures based on existing off-policy evaluation (OPE) methods and conduct extensive experiments in various environments to demonstrate the effectiveness of the proposed methodologies. In theory, we prove the optimality of the proposed treatment allocation design and establish upper bounds for the mean squared errors of the resulting treatment effect estimators.
Motivation. Prior to the full-scale deployment of any product in practical applications, an accurate evaluation of its potential impact is crucial. Modern technology companies, including Amazon, Google, Netflix, Uber, and Didi, commonly employ online experimentation or A/B testing as a means of evaluating the effectiveness of new products or policies (treatment) in comparison to their existing counterparts (control). Often, in these experiments, policies are assigned sequentially over time, impacting both current and future responses. Such a dynamic can invalidate the Stable Unit Treatment Value Assumption (SUTVA), as outlined by (Imbens and Rubin, 2015). Failure to consider these temporal carryover effects can lead to a biased estimation of the treatment effect. Additionally, the limited duration of the experiment and the typically small size of the treatment effect pose significant challenges to consistently detecting the treatment effect.
A substantial body of literature has been dedicated to developing policy evaluation or A/B testing algorithms using data from online experiments. However, there has been limited focus on the generation of the experimental dataset itself. This factor is critical as it can substantially influence the precision of the subsequent treatment effect estimator. This paper's primary focus is to study optimal experimental designs in the context of sequential decision making. In clinical trials, a carefully designed experiment can significantly improve the accuracy of the treatment effect estimator and the statistical power of the associated test, as noted by (Sverdlov et al., 2020). However, most existing
*The first two authors contribute equally to this paper. Address for correspondence: Hongtu Zhu, Ph.D., E-mail: htzhu@email.unc.edu.
Figure 1: Data structure of NMDP, TMDP, and MDP. (a) In the NMDP, the reward and future observations are determined by all past observation-action pairs. (b) In the TMDP, the reward and future observations depend solely on the current observation-action pairs. (c) In the MDP, the reward and future observations also rely on the current observation-action pairs, with the same-colored slashes indicating identical conditional distributions.
designs consider the classic identically and independently distributed (i.i.d.) setting, failing to account for temporal carryover effects.
Contributions. We make three important contributions below. Firstly, we propose three optimal dynamic treatment allocation strategies in sequential decision making, while accounting for carryover effects over time. The proposed designs are built upon the  $A$ -optimality criterion developed in non-dynamic settings (Atkinson et al., 2007), and effectively minimize the variance of the average treatment effect estimator when data follow a (time-invariant) Markov decision process (MDP), a time-varying MDP (TMDP) or a non-Markov decision process (NMDP). Figure 1 offers a visual illustration of these data generating processes. To our knowledge, this represents the first piece of work to provide an analytical form of optimal allocation strategies within these dynamic contexts.
Secondly, by leveraging insights from the established field of off-policy evaluation (OPE, see, for instance, Dudík et al., 2014; Uehara et al., 2022, for reviews), we devise subsequent estimation methodologies that learn the treatment effect estimator using log data produced in accordance with our proposed design. We employ real datasets gathered from a globally recognized ride-sharing company to create a city-scale synthetic environment that mimics the spatio-temporal dynamics of the order dispatching problem. Our findings indicate that the proposed estimators deliver significantly superior accuracy compared to conventional baseline estimators.
Thirdly, we establish the theoretical properties of the proposed designs and estimators. Specifically, we establish that in NMDPs, the optimal design distributes the initial policy randomly, with probabilities corresponding to the standard deviations of the cumulative rewards, and subsequently implements the same policy. However, such a design may not necessarily be optimal in (T)MDPs. To address this, we focus on a category of restricted designs and identify sufficient conditions that ascertain their asymptotic optimality in (T)MDPs. Moreover, we derive an upper limit for the mean squared error (MSE) of the resulting treatment effect estimator. Extensive simulation studies have been conducted to assess the finite sample performance of the proposed algorithms. Python code implementing the proposed algorithms is available at https://github.com/tingstat/MDP/design.
Experimental designs. There is an extensive body of literature on experimental design for clinical trials, with a multitude of optimal designs proposed. These include the global treatment balance design (Efron, 1971; Wei, 1977), designs that regulate the marginal covariate distribution across various treatment groups (Pocock and Simon, 1975; Heritier et al., 2005; Liu and Hu, 2022), and criteria-based design anchored on variance minimization (Begg and Iglewicz, 1980; Wong and Zhu, 2008; Chandereng et al., 2020). Other noteworthy designs are the  $D$ -optimality and its generalization  $D_A$  optimality (Jones and Goos, 2009; Atkinson and Pedrosa, 2017; Loux, 2013; Liu et al., 2022), alongside the  $A$ -optimality and its extension  $A_A$  optimality (Atkinson et al., 2007; Sverdlov and Rosenberger, 2013; Yin and Zhou, 2017). Moreover, numerous sequential adaptive designs for clinical trials have been developed, including covariate-adaptive designs (Zhu and Hu, 2019; Wang
and Ma, 2021), response-adaptive designs (Zhu et al., 2020; Yu et al., 2022), and covariate-adjusted response-adaptive designs (Zhang et al., 2007; Villar and Rosenberger, 2018). However, these methods were designed for i.i.d. data and thus are not directly applicable to our settings.
Recently, several papers have studied experimental designs with spatial/network spillover effects (Ugander et al., 2013; Baird et al., 2018; Li et al., 2019; Jagadeesan et al., 2020; Nandy et al., 2020; Viviano, 2020; Zhou et al., 2020; Kong et al., 2021; Leung, 2022). Additionally, a few designs have been developed for modern technological companies, such as e-commerce (Bajari et al., 2021; Nandy et al., 2021) and ride-sharing (Johari et al., 2022). However, these studies differ from ours in that they did not utilize NMDP or TMDP models for experimental designs.
Bojinov et al. (2023) explored optimal regular switchback design in the presence of a single experimental unit. Their approach relies on the assumption that the carryover effects last for a fixed duration. However, this assumption is not consistent with our NMDP and (T)MDP models, where the carryover effects, due to state transitions, can potentially last indefinitely.
Finally, Bayesian Optimal Experimental Design (BOED) presents a powerful tool for numerically determining the optimal design (see e.g., Ryan et al., 2016; Rainforth et al., 2023). Recently, works such as Foster et al. (2021), Lim et al. (2022), and Blau et al. (2022) have proposed leveraging deep reinforcement learning to compute the BOED. In contrast, our paper employs a frequentist approach, analytically deriving optimal designs for sequential decision making.
Off-policy evaluation. Our work is closely related to OPE, which seeks to estimate the mean outcome of a new target policy using logged data collected by a different policy. It has been applied in a range of domains including healthcare (Murphy, 2003), education (Mandel et al., 2014), ride-sharing (Shi et al., 2022) and video-sharing (Xu et al., 2023). Recently, it has been employed to conduct A/B testing and mediation analysis in the presence of temporal carryover effects (Hu and Wager, 2022; Farias et al., 2022; Shi et al., 2022; Tang et al., 2022; Ge et al., 2023; Shi et al., 2023). Existing algorithms include value-based methods (Mannor et al., 2007; Antos et al., 2008; Le et al., 2019; Feng et al., 2020; Luckett et al., 2020; Hao et al., 2021; Liao et al., 2021; Chen and Qi, 2022; Shi et al., 2022; Bian et al., 2023), importance (re)sampling-based methods (Swaminathan and Joachims, 2015; Liu et al., 2018; Schlegel et al., 2019; Yin and Wang, 2020; Thams et al., 2021; Wang et al., 2021; Hu and Wager, 2023), and doubly robust or more robust methods (Zhang et al., 2013; Jiang and Li, 2016; Thomas and Brunskill, 2016; Farajtabar et al., 2018; Kallus and Uehara, 2020; Uehara et al., 2020; Shi et al., 2021; Kallus and Uehara, 2022; Liao et al., 2022; Xie et al., 2023). In particular, Kallus and Uehara (2022) and Liao et al. (2022) established the efficiency bounds for OPE, which we leverage to construct optimal designs. However, none of the previous work considered how to design the behavior policy that generates high-quality data to improve the efficiency of the policy value estimator.
Recently, Wan et al. (2022) studied safe experimental designs for efficient policy evaluation under certain safety constraints in a non-dynamic setting without temporal carryover effects. However, it remains unclear how to generalize their methodologies to sequential decision making.
A/B testing. Lastly, our proposal aligns with a body of literature that develops A/B testing algorithms for randomized online experiments (see e.g., Kharitonov et al., 2015; Johari et al., 2017; Yang et al., 2017; Bojinov and Shephard, 2019; Ju et al., 2019; Zhou et al., 2020; Shi et al., 2021; Maharaj et al., 2023; Wan et al., 2023). However, most works focus on the sequential monitoring problem, which involves partitioning the significance level at each interim stage to control the overall type-I error. This is distinct from our goal, which is to generate high-quality experimental data to enhance the efficiency of estimation and inference.
2 Models and Problem Formulation
Data. We consider an experiment spanning  $n$  days and each day is divided into  $T$  time intervals. On the  $i$ -th day  $(i = 1, \dots, n)$  and time interval  $t$ $(t = 1, \dots, T)$ , the decision maker observes certain time-varying market features, denoted as  $O_{t}^{(i)}$  (e.g., the numbers of incoming orders and available drivers in a ride-sharing platform) and chooses to implement one of the two policies. This action is denoted as  $A_{t}^{(i)} \in \{0, 1\}$ . By convention,  $A_{t}^{(i)} = 1$  indicates the implementation of a new policy, while  $A_{t}^{(i)} = 0$  signifies the use of the control policy. Afterwards, the decision maker receives an immediate reward  $R_{t}^{(i)}$  and subsequently observes the next observation  $O_{t+1}^{(i)}$ . This process is
repeated until we reach the termination time  $T$ . The observed data can thus be summarized as  $\{O_{t}^{(i)}, A_{t}^{(i)}, R_{t}^{(i)}, t = 1, \ldots, T\}_{i=1}^{n}$ .
Model. We consider three models for the aforementioned data generating process: an NMDP, a TMDP, and a (time-invariant) MDP. As depicted in Figure 1, an NMDP is a versatile model that doesn't impose Markovian or stationarity assumptions. The immediate reward  $(R_{t})$  and future observation  $(O_{t + 1})$  depend on the entire history of observations and actions, not merely the current observation-action pair  $(O_{t}$  and  $A_{t})$ . This kind of model has been extensively used in research on learning optimal dynamic treatment regimes (Kosorok and Laber, 2019; Tsiatis et al., 2019). Conversely, a TMDP is a special case of NMDP, with a stronger structural assumption that satisfies Markovanity. It assumes that the current observation-action pair is sufficient to determine the immediate reward and future observation, while allowing this conditional distribution function to vary with time. This type of data generating process is also known as an episodic MDP and has been widely studied in the reinforcement learning literature (see e.g., Jin et al., 2018, 2021; Li et al., 2023). Finally, an MDP is a special TMDP with an additional stationarity assumption.
Objective. A history-dependent policy, denoted as  $\pi$ , is a sequence of decision rules  $\{\pi\}_{t \geq 1}$ . Each rule,  $\pi_t(\bullet | H_t)$ , uses the observed data history  $H_t$  up to time  $t$  as input and provides a probability distribution as output. This output determines the likelihood of choosing each action  $A_t$  at time  $t$ . The primary objective here is to estimate the Average Treatment Effect (ATE), which is defined as:
$$
\mathrm {A T E} = T ^ {- 1} \sum_ {t = 1} ^ {T} \left[ \mathbb {E} ^ {1} \left(R _ {t}\right) - \mathbb {E} ^ {0} \left(R _ {t}\right) \right],
$$
where  $\mathbb{E}^1$  (or  $\mathbb{E}^0$ ) signifies the expected value when the system follows the new (or control) policy. Moreover, for any  $a$  in the set  $\{0,1\}$ , let's define  $V_{t}^{a}(h)$  as the value function,  $\sum_{k = t}^{T}\mathbb{E}^{a}(R_{k}|H_{t} = h)$ , when a certain action is followed. It's evident from this that the ATE equals  $T^{-1}\mathbb{E}[V_1^1 (O_1) - V_1^0 (O_1)]$ . Let  $\pi^b$  denote the behavior policy that generates the experimental data, i.e.,  $\pi_t^b (a|H_t) = \mathbb{P}(A_t = a|H_t)$  for any  $a$  and  $t$ . Our objective lies in the design of an optimal behavior policy to generate high-quality data so that the mean squared error of the subsequent ATE estimator is minimized.
3 Design, Implementation and Evaluation in NMDPs
We focus on NDMPs in this section. We begin by proposing the optimal dynamic treatment allocation strategy to generate the experimental data. We next discuss some implementation details to implement the proposed design. Finally, we construct the ATE estimator from the data collected.
Design. Jiang and Li (2016) and Kallus and Uehara (2020) established the semiparametric efficiency bound, which is a nonparametric extension of the Cramer-Rao lower bound in parametric models (Bickel et al., 1993), for offline estimation of the average return under a general target policy in NMDPs. In essence, the efficiency bound corresponds to the smallest mean squared error (MSE) among a broad class of regular off-policy estimators. They further developed semiparametrically efficient estimators whose asymptotic MSEs reach this lower bound. Based on their results, one can easily show that the efficiency bound (EB) for the ATE equals:
$$
\mathbf {E B} _ {1} \left(\pi^ {b}\right) = T ^ {- 2} \sum_ {a \in \{0, 1 \}} \sum_ {t = 1} ^ {T} \mathbb {E} ^ {\pi^ {b}} \left[ \sigma_ {t} \left(H _ {t}, a\right) \prod_ {k \leq t} \frac {\mathbb {I} \left(A _ {k} = a\right)}{\pi_ {k} ^ {b} \left(a \mid H _ {k}\right)} \right] ^ {2} + T ^ {- 2} \operatorname {V a r} \left[ V _ {1} ^ {1} \left(O _ {1}\right) - V _ {1} ^ {0} \left(O _ {1}\right) \right], \tag {1}
$$
where  $\sigma_t^2 (H_k,a)$  denotes the conditional variance of the temporal difference error  $R_{t} + V_{t + 1}^{a}(H_{t + 1}) - V_{t}^{a}(H_{t})$  given  $H_{t}$  and that  $A_{t} = a$ , and  $\mathbb{I}(\bullet)$  denotes the indicator function. Note that the second term on the right-hand-side (RHS) is independent of the behavior policy. However, the first term is a function of  $\pi^b$  in that: (i) the ratio in the first term explicitly involves  $\pi^b$ ; (ii) the expectation  $\mathbb{E}^{\pi^b}$  is taken with respect to the data generated by  $\pi^b$ .
The proposed design identifies the optimal  $\pi^{b*}$  that minimizes (1). Before delving into the details of the proposed design, we first present the main idea behind it. A key observation is that the cumulative ratio  $\prod_{k\leq t}[\mathbb{I}(A_k = a) / \pi_k^b (a|H_k)]$  appears on the RHS of (1) due to the use of sequential importance sampling (see e.g., Jiang and Li, 2016, Equation (4)) to account for the distributional shift between  $\pi^b$  and the global policy that assigns action  $a$  at each time. In general, the value of the cumulative ratio grows with the discrepancy between the target and behavior policy. Hence, it is plausible to
Algorithm 1 Treatment allocation algorithm for NMDPs
Input: The burn-in period  $m_0$  for each global policy and the termination day  $n$ .
1: Run each global policy for  $m_0$  days and obtain  $\{O_t^{(i)},A_t^{(i)},R_t^{(i)}\}_{i = 1}^{2m_0}$
2: while  $2m_0 < m \leq n$  do
3: Using the collected data to estimate  $V_{1}^{a}$  and  $\sigma_{*}(\bullet, a)$  via (3). Obtain  $\widehat{\sigma}_{*}(O_{1}^{(m)}, a)$ .
4: Assign  $A_1^m$  according to  $\widehat{\pi}_{1,m-1}^{b*}(a|O_1^{(m)}) = \widehat{\sigma}_*(O_1^{(m)}, a) / [\widehat{\sigma}_*(O_1^{(m)}, 1) + \widehat{\sigma}_*(O_1^{(m)}, 0)]$ .
5: Set  $A_2^{(m)} = \dots = A_T^{(m)} = A_1^{(m)}$ .
6: end while
Output:  $\{O_t^{(i)},A_t^{(i)},R_t^{(i)}\}_{i = 1}^n$
expect that an on-policy estimator, where all actions are determined according to the target policy, could potentially minimize the MSE.
We show in Theorem 1 below that the optimal design corresponds to a slight modification of the aforementioned on-policy design. Specifically, it randomly allocates the initial action with probabilities proportional to the standard deviations of the cumulative rewards (see Equation (2)), and assigns the same action afterwards.
Theorem 1. In NMDPs,  $\pi^{b*}$  that minimizes (1) satisfies: (i) for any  $a \in \{0,1\}$
$$
\pi_ {1} ^ {b *} (a | O _ {1}) = \frac {\sigma_ {*} (O _ {1} , a)}{\sigma_ {*} (O _ {1} , 0) + \sigma_ {*} (O _ {1} , 1)}, \tag {2}
$$
where  $\sigma_{*}^{2}(O_{1},a) = \mathbb{E}^{a}[\{\sum_{t}(R_{t} - \mathbb{E}^{a}R_{t})\}^{2}|O_{1},A_{1} = a]$ ; (ii)  $\pi_2^{b*}(A_1|H_2) = \pi_3^{b*}(A_1|H_3) = \dots = \pi_T^{b*}(A_1|H_T) = 1$  almost surely, or equivalently,  $A_{1} = A_{2} = \dots = A_{T}$  under  $\pi^{b*}$ .
Implementation. It remains to estimate the standard deviation  $\sigma_{*}$  to implement the proposed design. If historical data is available, we can use it to estimate  $\sigma_{*}$ . Otherwise, we can use data collected from the experiment to adaptively update this parameter. Initially, we run each global policy for  $m_0$  days to generate the data. Next, on the  $m$ th day ( $2m_0 < m \leq n$ ), we estimate  $\sigma_{*}$  using the experimental data collected so far. To this end, we begin by applying existing supervised learning algorithm to estimate the value function  $V_1^a$  by regressing the cumulative rewards  $\{\sum_t R_t^{(i)} : A_1^{(i)} = a, i < m\}$  on the initial observations  $\{O_1^{(i)} : A_1^{(i)} = a, i < m\}$  where the superscript  $i$  indicates that the data are collected on the  $i$ th day. Let  $\widehat{V}_{1,m-1}^a$  denote the resulting estimator. We next employ supervised learning again to regress the squared residuals on the initial observations to estimate  $\sigma_{*}(\bullet, a)$  as
$$
\widehat {\sigma} _ {*} (\bullet , a) = \arg \min  _ {\sigma (\cdot)} \sum_ {i = 1} ^ {m - 1} \mathbb {I} (A _ {1} ^ {(i)} = a) \left\{\left[ \sum_ {t} R _ {t} ^ {(i)} - \widehat {V} _ {1, m - 1} ^ {(a)} (O _ {1} ^ {(i)}) \right] ^ {2} - \sigma^ {2} (O _ {1} ^ {(i)}) \right\} ^ {2}. \tag {3}
$$
Finally, we plug  $\widehat{\sigma}_{*}(O_{1}^{(m)},a)$  into the RHS of (2) to obtain  $\widehat{\pi}_{1,m - 1}^{b*}$ , assign  $A_{1}^{(m)}$  according to  $\widehat{\pi}_{1,m - 1}^{b*}$  and set  $A_{2}^{(m)} = \dots = A_{T}^{(m)} = A_{1}^{(m)}$ . We summarize our procedure in Algorithm 1.
Evaluation. Finally, we compute the ATE estimator using the experimental data. Several OPE algorithms, including value-based, importance sampling (IS), and doubly robust (DR) methods, are applicable for this purpose. DR estimators are known for achieving the efficiency bound (Kallus and Uehara, 2020). In our context, we suggest using the following online DR estimator for ATE,
$$
\widehat {\mathrm {A T E}} _ {1} = \sum_ {a = 0} ^ {1} \frac {(- 1) ^ {a + 1}}{T (n - 2 m _ {0})} \sum_ {i = 2 m _ {0} + 1} ^ {n} \left[ \widehat {V} _ {1, i - 1} ^ {a} \left(O _ {1} ^ {(i)}\right) + \frac {\mathbb {I} \left(A _ {1} ^ {(i)} = a\right)}{\widehat {\pi} _ {1 , i - 1} ^ {b *} \left(a \mid O _ {1} ^ {(i)}\right)} \left[ \sum_ {t} R _ {t} ^ {(i)} - \widehat {V} _ {1, i - 1} ^ {a} \left(O _ {1} ^ {(i)}\right) \right] \right]. \tag {4}
$$
Note that both the estimated value function  $\widehat{V}_{1,i-1}^{a}$  and behavior policy  $\widehat{\pi}_{1,i-1}^{b*}$  in (4) are computed during the data collection process. These nuisance functions are independent of the data  $\{O_{1}^{(i)}, A_{1}^{(i)}, R_{t}^{(i)}\}$  used to construct the policy value. This cross-fitting procedure enables us to circumvent the need to impose certain metric entropy conditions (Díaz, 2020) and has been widely employed in the statistics and machine learning literature (Luedtke and Van Der Laan, 2016; Bibaut et al., 2021; Shi et al., 2021). By design, (4) can be calculated in an online manner, eliminating the need to store historical data. Moreover, we implement a burn-in procedure that discards the first  $2m_{0}$
samples to construct  $\widehat{\mathsf{ATE}}$ . This ensures the consistencies of  $\widehat{V}_{1,i}^{a}$  and  $\widehat{\pi}_{1,i}^{b*}$ . Concurrently, value-based and IS-based methods are also suitable for estimating ATE. Provided that the nuisance functions are properly modeled, these estimators can achieve the efficiency bound as well (Uehara et al., 2020; Liao et al., 2021; Wang et al., 2021; Shi et al., 2023). The following theorem provides an upper bound for the MSE of the proposed ATE estimator.
Theorem 2. Suppose that  $\min_{i}\widehat{\pi}_{1,i}^{b*}\geq \epsilon$  and  $V_{1,i}^{a}\leq TR_{\max}$  for some constants  $\epsilon >0$  and  $R_{\max} < \infty$ ,  $\max_{a,i}\mathbb{E}|1 / \widehat{\pi}_{1,i}^{b*}(a,O_1) - 1 / \pi_1^{b*}(a,O_1)|_2^2\leq Ci^{-2\alpha_1}$  and  $\max_{a,i}T^{-2}\mathbb{E}|\widehat{V}_{1,i}^{a}(O_{1}) - V_{1,i}^{a}(O_{1})|_{2}^{2}\leq$ $C i^{-2\alpha_{2}}$  for some constants  $C < \infty,0 < \alpha_{1},\alpha_{2} < 1 / 2$ . Then we have
$$
\mathbb {E} \left(\widehat {A T E} _ {1} - A T E\right) ^ {2} \leq \frac {E B _ {1} \left(\pi^ {b *}\right)}{n - 2 m _ {0}} + O \left(\epsilon^ {- 1} C (n - 2 m _ {0}) ^ {- 1 - 2 \alpha_ {2}}\right) + O \left(R _ {\max } ^ {2} \sqrt {C} (n - 2 m _ {0}) ^ {- 1 - \alpha_ {1}}\right).
$$