Stage-Aware Learning for Dynamic Treatments

Consider a balanced multistage decision setting where all patients in the study have a total number of T stages (visits). For each patient at their $j^{th}$ clinic visit, where $j=1,\ldots,T$, a set of time-varying variables $X_{j}\in\mathcal{X}_{j}$ are recorded to collect individual health status. Consequently, a new treatment assignment $A_{j}$ is delivered based upon the patient’s longitudinal historical information from their first visit to the $j^{th}$ visit, denoted by $H_{j}=S_{j}(X_{1},A_{1},X_{2},A_{2},\ldots,A_{j-1},X_{j})\in\mathcal{H}_{j}$, where $S_{j}$ is a deterministic summary function. In this paper, we consider binary treatments, i.e., $A_{j}\in\mathcal{A}=\{1,-1\}$, where the interpretation of the treatment options are nested depending on the treatments assigned in previous stages. After the final visit, a clinical outcome $R=\sum_{j=1}^{T}r_{j}$, also known as the total reward from all immediate rewards, is obtained to reflect the benefits of the allocated treatment assignment.

A DTR is a sequence of decision rules $\mathscr{D}=\{D_{j}:\mathcal{H}_{j}\mapsto\mathcal{A}\}_{j=1}^{T}$ that map patients’ historical information onto treatment space. The decision rules $\mathscr{D}$ could also be represented by a composite function of the real-valued functions $\mathbf{f}=\{f_{j}\in\mathcal{F}:\mathcal{H}_{j}\mapsto\mathbb{R}\}_{j=1}^{T}$, where a realization of the decision that follows the treatment regime at the $j^{th}$ visit is $d_{j}=D_{j}(H_{j})=\text{sign}\{f_{j}(H_{j})\}$. Now, assume that the full trajectory of an observation sequence $\{X_{1},A_{1},X_{2},A_{2},\ldots,X_{T},A_{T},R\}$ follows a data distribution $P$. Our goal is to seek the optimal treatment regime $\mathscr{D}^{*}$ which yields the largest expected rewards among all regimes:

Note that the expectation operator $\mathbb{E}^{\mathscr{D}}$ in equation (1) is taken with respect to an unknown restricted distribution $\{X_{1},A_{1}=d_{1},X_{2},A_{2}=d_{2},...,X_{t},A_{T}=d_{T},R\}\sim P^{\mathscr{D}}$, which describes the probability distribution when the treatments are assigned according to the regime $\mathscr{D}$. By convention, we call the corresponding $\mathscr{D}$ the target regime. Since the historical information and potential outcome under $P^{\mathscr{D}}$ are unobservable, to infer the decision rules from the observed data while avoiding confounding issues between the assignments and expected rewards, we adopt the following three standard assumptions:

Under the aforementioned assumptions, it can be shown that the expected total reward under the target regime $\mathscr{D}$ is estimable by inverse probability weighting (Qian and Murphy, 2011):

where $\pi_{j}$ is the propensity score function and the density corrected expected reward is referred to as the IPWE. Provided that the propensity functions are correctly specified, the optimal treatment regime $\mathscr{D}^{*}$ is the maximizer of the IPWE. Accordingly, learning schemes that directly estimate the optimal regime by maximizing the IPWE estimator could be categorized as IPWE-based approaches (Zhao et al., 2012, 2015; Laha et al., 2024).

However, due to the density ratio corrections, IPWE-bsaed approaches only incorporate a reward when an individual’s treatments are fully matched with the target regime. Such a full-matching requirement can be reflected in the conditional expectation form of the weigthed total rewards,

which is determined by two major factors: first, the expected reward among the population whose observed treatments are fully matched with the target regime $\mathscr{D}$, i.e., full-matching expected reward; and second, the probability of assigning treatments conforming with the target regime at all stages, i.e., target regime assignment rate. At the population level, the probability of assigning any arbitrary dynamic regime is ensured to be non-zero due to the positivity assumption. However, if the optimal regime assignment rate is small among the patient population, it is highly possible that none of the sampled patients may fully follow the optimal treatments at all stages, and thus the optimal regime is infeasible in practice. Clearly, this strict full-matching requirement may lead to difficulties on both optimization and estimation. We call this phenomenon the curse of full-matching.

To fully understand the dilemma of full-matching, we frame this challenge within the context of a randomized trial and provide a concrete illustrative example presented in Table 1. Similar to observational studies, where the behavior policy can rarely keep making optimal decisions at all stages, this example assigns the optimal regime, Treatment Arm 1, at a small rate of 0.15 (i.e., $P(A_{1})\cdot P(A_{2}|A_{1})=0.5\cdot 0.3=0.15)$. Due to the requirement of full-matching, IPWEs only focus on identifying these 15% of subjects assigned to Treatment Arm 1 during empirical estimations. As a result, it might require many more samples to be collected before IPWE-based approaches could estimate the optimal regime and achieve near-asymptotic properties. Furthermore, note that the above example consists of only two decision stages. In practice, a growing number of stages tends to make estimation even harder since the target regime assignment rate decreases exponentially.

Among existing literature, methods are developed to stabilize IPWEs from this strict full-matching condition. For instance, the robust estimators (Zhang et al., 2012a, b, 2013; Zhao et al., 2019; Schulz and Moodie, 2021) augment IPWEs with an unbiased term to achieve double-robustness. This provides a safeguard when estimating the propensity of the optimal regime becomes unstable due to the curse of full-matching. RWL (Zhou et al., 2017) estimates outcome residuals to stabilize IPWEs against outcome shifts. Q-learning based policy searches, such as C-learning (Zhang and Zhang, 2018), leverage the advantages of regression-based frameworks, which do not rely on probability weighting and circumvent the curse of full-matching. However, these methods all require accurate outcome models for their additional augmentation terms or regression components. Given the high heterogeneity among rewards and an increasing number of decision stages, correctly specifying the outcome model poses significant challenges and increases computational complexity.

Our work seeks to address the challenges posed by the curse of full matching while maintaining the simplicity of IPWEs and avoiding the specification issues commonly associated with outcome models. Specifically, referring back to the example in Table 1, we aim to leverage the substantial information available in Treatment Arms 2 and 3, which differ from the optimal regime by only one treatment but are more prevalent in the trial with a combined probability of 70% (i.e., $P(\text{arm}2\text{ or }\text{arm}3)=0.5\cdot(0.7+0.7)=0.7$). Notably, the total rewards from these non-optimal treatment arms, especially Treatment Arm 3, are not significantly different from the optimal arm. Inspired by this observation, we propose a novel method which can incorporate these non-optimal treatments in the estimation of the optimal regime and break the curse of cull-matching to improve sample efficiency.

The performance of IPWE-based approaches is largely restricted by the full-matching assignment rate of the optimal regime among the patient populations. Though we might not control how treatments are administrated according to the optimal regime at every decision stage, in this subsection, we propose to relax the strict full-matching requirement by allowing decision discrepancies between the assigned treatments and the target regime $\mathscr{D}$ at $T-k$ number of stages, where $0\leq k\leq T$. In other words, the optimal regime is allowed to partially match the treatment sequence at exactly $k$ number of arbitrary decision stages. The rationale is that, while it is unrealistic to expect all decisions to be optimal, encountering one or more non-optimal decisions at some stages is more likely in practice. As a result, the probability of patients receiving optimal decisions at $k$ number of stages could be much larger than the probability of patients receiving optimal decisions at all stages. We call this relaxed version of the full-matching requirement the k-partially matching requirement.

To formalize the notation, we let random variable $K$ denote the number of correct alignments between treatments $\mathbf{A}=\{A_{j}\}_{j=1}^{T}$ and decisions from an arbitrary target regime,

At each stage, we examine whether the assigned treatment aligns with the recommended treatment. Importantly, the recommended treatment is determined based on the patient’s historical information, which makes treatments nested and their effects carried over the decision stages. Additionally, if we specify $K$ to be a realized value $k$ within the range $\{0,\ldots,T\}$, we are constraining our optimal regime search to the patient population who are $k$-partially matched with the optimal regime. A more extreme $k$ value (e.g., $k=0$ or $k=T$) corresponds to a more restricted alignment requirement. Under the IPWE framework, only patients with $K=T$ are included in the estimation procedure.

When $K=k$, a new restricted unknown distribution is induced, i.e., $(X_{1},A_{1}=\tilde{d}_{1},...,X_{T},A_{T}=\tilde{d}_{T},R)\sim P^{\mathscr{D}_{(k)}}$, where $\tilde{d}_{j}=(-1)^{\mathbb{I}(j\in\mathcal{K})+1}\cdot d_{j}$ indicates that at any stage $0\leq j\leq T$, the decision $\tilde{d}_{j}$ is the same as the target regime $d_{j}$ only if $j$ is among the indexes of $k$ arbitrary matching stages $\mathcal{K}$ (i.e., $\tilde{d}_{j}=d_{j}\;\text{if}\;j\in\mathcal{K}$). Subsequently, $P^{\mathscr{D}_{(k)}}$ is a distribution of an observation sequence where its $k$ out of T assignments $\{A_{j}\}_{j\in\mathcal{K}}$ are followed by the regime $\mathscr{D}$. Based on the derivation of $k$-matching potential outcomes provided in Appendix A.1, we can similarly adopt the density ratio correction and obtain an IPWE for the expected rewards evaluated under the new measure $P^{\mathscr{D}_{(k)}}$. We denote the new estimator $\mathbb{E}^{\mathscr{D}_{(k)}}[R]$ as $k$-IPWE and present the results in Proposition 6:

The regime $\tilde{\mathscr{D}}_{(k)}$ maximizing the $k$-IPWE would yield the largest expected reward if patients were treated by $\mathscr{D}$ at $k$ number of stages. In addition, note that we do not require the $k$ matching stages to be the same for each individual. As long as there is an exact $k$ number of treatment matchings between the assignments and target regime $\mathscr{D}$, those patients’ rewards are involved in the maximization process. As a result, the $k$-IPWE provides a superior level of flexibility.

However, the performance of the proposed $k$-IPWEs still depends on the pre-selected $k$ value. The purpose of designing the $K$ treatment matching number is to increase the conditional probability of the constrained population receiving optimal treatments. When the value $k$ is poorly selected, the probability of patients being $k$-partially matched could be small, and we could encounter a similar aforementioned empirical dilemma. For instance, in a population with a 99% full-matching assignment rate, specifying the random variable $K$ with values other than $T$ leads to a small $k$-partially matching rate. In other words, the curse of full-matching can be effectively minimized only if the pre-specified $k$ has the highest $k$-partially matching probability, i.e., $k=\operatorname*{argmax}_{k\in\{0,\ldots,T\}}P(K=k)$. Finding such a $k$ value is possible but computationally cumbersome, and yet not every patient will participate in the optimization process due to the population conditional constraints. To reduce the uncertainty of selecting $k$ values and include all individuals from the sample, we further construct a learning method based on $k$-IPWE which can incorporate all scenarios of $k$-partially matching through applying a weighting scale on the matching number $K$.

In this subsection, we formally introduce a novel learning method to combine all levels of $k$-IPWEs into one single estimation task. This addresses the dependencies of $k$-IPWEs at the pre-selected $k$-values. Since the new estimator accounts for treatment and regime matching status at any number of stages, we name the new learning method Stage-Aware Learning (SAL).

To start with, we re-weight each $k$-IPWE by $k/T$, proportional to the number of matching stages $k$, and estimate the optimal regime simultaneously by maximizing the following SAL value function derived in Appendix A.2,

We denote the estimated maximizing regime $\tilde{\mathscr{D}}=\operatorname*{argmax}_{\mathscr{D}}V^{SA}(\mathscr{D})$. In our choice, the applied weights increase with the $k$ value, and imply that we prioritize the regime which has closer alignment with the optimal decisions while achieving higher expected rewards during estimation. Structurally, the weighting component of the final SAL value function resembles the formulation of the Hamming loss (Tsoumakas and Katakis, 2007), which has the following three unique advantages.

First of all, compared to the IPWE value function (2), the SAL value function replaces the product of indicator functions with the correct treatment alignment percentage. Therefore, instead of excluding patients completely if one of their observed treatments is not aligned with the optimal decision, SAL still considers those patients but discounts their rewards based on the degree of alignment between observed and optimal treatments across all of the decision stages. That is, even if the decision sequence is long, the new learning process is able to utilize all available patients’ outcomes and maximize the alignment percentages over those with high rewards. As a result, all patients and their treatment strategies are included in the optimal regime searching procedure. Second, the SAL learning scheme is analogous to a multi-label classification framework. Instead of employing a surrogate function involving all decision stages simultaneously in OWL (Zhao et al., 2015), our proposed regime can be optimized at each individual stage to match the optimal decisions. This leads to computational convenience as the optimization of the Hamming loss function has been well established and the multi-stage learning task can be segmented into single-stage sub-tasks.

In fact, SAL suggests a more general DTR learning framework by allowing a probability distribution on the matching number $K$. Particularly, the IPWE-based framework is a special case of ours under the general framework indicated in Remark 2. Suppose the density function of $K$ is proportional to a scale function $\phi(\cdot)$. After taking the iterated expectation of $k$-IPWEs under all possible choices of $k$ values, we obtain a new estimator of the expected rewards aimed to maximized for any arbitrary target regime $\mathscr{D}$,

The new estimator results in a generic form where the total rewards are weighted by the density of $K$. Correspondingly, the maximization procedure not only finds the regime that yields the largest expected total rewards, but also identifies the one which matches with the observed treatment sequence closest to the underlying distribution of the treatment matching number. In addition, compared to the previous regime estimator $\tilde{\mathscr{D}}_{(k)}$ based on the $k$-IPWE in equation (6), the general estimator no longer constrains the rewards under a sub-population of patients. Instead, it takes every patient’s reward into account as long as the density is positive for all matching number $K$, i.e., $\phi(k)>0$ $\forall k\in\{1,..,T\}$.

In summary, we propose a novel SAL method under a more general DTR estimation framework, to improve data efficiency and empirical adaptivity of IPWE-based methods. By combining each level of $k$-IPWEs, we include all matching scenarios and relax the selection of the $K$ value. Through imposing higher weights to $k$-IPWEs with larger $k$ value, SAL possesses the interpretation of searching samples with high total rewards and large optimal treatment matching percentages simultaneously without the need for stage immediate rewards. However, the absence of immediate rewards may complicate the learning process of treatment effects at each stage. It requires larger sample sizes for SAL to attribute variations in total rewards to a single stage, especially when dealing with individuals having similar matching percentages. Next, we further propose a weighted learning scheme based on SAL to incorporate stage heterogeneity and facilitate the DTR learning process.

In this subsection, we propose a non-trivial variant of SAL, namely the Stage Weighted Learning (SWL) method, based on a weighted multi-label framework to enhance stage heterogeneity and facilitate distributing stage-wise treatment effects from total rewards into individual stages, We formulate the SWL value function as follows,

where $\{\omega_{j}\}_{j=1}^{T}$ are the real-valued weights satisfying $\omega_{j}\in[0,1]$ and $\sum_{j=1}^{T}\omega_{j}=1$ for any $j$ in $\{1,..,T\}$. Intuitively, a larger weight is imposed on a stage with more substantial treatment effect contributing to the total rewards. As the weights are designed to quantify the relative importance of treatment effects among decision stages, we denote the imposed weights $\{\omega_{j}\}_{j=1}^{T}$ as stage importance scores.

With the incorporation of stage importance scores, a notable advantage of SWL is the rescaling of the reward by a weighted average of the treatment alignments. Compared to the SAL’s uniform stage weights ($1/T$) applied to each treatment-matching stage outlined in equation (7), the nonidentical stage importance scores introduce stage heterogeneity to the matching percentage component of the value function. Specifically, individuals with the same number of alignments between assigned treatments and optimal decisions no longer have identical treatment-matching percentages as seen in SAL. Instead, their matching percentages vary based on the importance scores assigned to the matched stages, i.e., a higher matching percentage is achieved when more stages with larger importance scores are matched. Consequently, in addition to the total rewards which capture the combined treatment effects across all decision stages, the matching percentages augment the stage-wise treatment effect through the importance scores, and thus further account for the stage heterogeneity within the SWL learning scheme.

The proposed stage importance scores also enhance the learning process of treatment regimes. As the total rewards are scaled by the weighted matching percentages in equation (9), treatment mismatch at stages with large importance scores would lead to a substantial loss in total expected rewards; whereas the expected rewards only have minor fluctuations at stages with negligible importance scores, regardless of treatment assignments. Under the main objective of maximizing the expected total rewards, SWL consequently prioritizes improving treatment alignment accuracy at important stages, and the importance scores at the individual stage level can effectively direct SWL’s attention towards those stages with substantial treatment effects.

Noticeably, one of the key components in SWL is the stage importance score. However, estimating these scores is non-trivial, mainly because the immediate rewards used to evaluate the treatment effects at each stage in some cases are unobservable. To solve this challenge, we utilize the attention-based mechanism (Bahdanau et al., 2014). Suppose there exists an underlying immediate reward function structure $\{\mathbf{\tilde{r}}_{j}\}_{j=1}^{T}$ which takes the stage importance score as an additional parameter. Formally, it affects individual’s immediate rewards at the $j^{th}$ stage as,

where $\mathbf{r}_{j}:\mathcal{H}\times\mathcal{A}\mapsto\mathbb{R}^{+}$ is the conventional definition of the immediate reward function. Under our assumption, the importance scores can be disentangled from the original rewards and are invariant to individual patients, representing stage-wise heterogeneity. Then, due to the fact that the expected total reward is the summation of expected immediate rewards, we estimate the importance scores scalars by minimizing the empirical squared loss between total rewards and constructed surrogate rewards from a semi-parametric point of view, i.e.,

Note that we do not limit the parametric form of the reward function space $\mathcal{R}$. For illustration purposes, we represent each reward function $\tilde{\mathbf{r}}_{j}$ with a fully-connected (FC) network due to its flexible capability of function approximation (DeVore et al., 2021), and adopt a long-short term memory (LSTM) network (Hochreiter and Schmidhuber, 1997) to capture the unobserved patients’ historical information $H_{ij}$ using up-to-date patients’ covariate information $\{X_{it}\}_{t=1}^{j}$ and past treatments $\{A_{it}\}_{t=1}^{j-1}$. Finally, we leverage the attention mechanism and propose an attention-based recurrent neural network architecture in Figure 1 to estimate the stage importance scores.

The idea behind the attention mechanism is to scale an input sequence by relevance to the predicting outcomes, with a more relevant part being assigned a higher weight. These weights focus the attention of the prediction model on the most relevant part of the input sequence to improve model performance. In our scenario, we view the stage importance scores as the attention weights so that stages with higher contributions and more relevance to the final total rewards possess larger attention weights. Therefore, the weights not only explicitly impose stage heterogeneity on the proposed neural network, but also direct the network to pay more attention to the stages with large treatment effects when predicting the total rewards. Once we compute the surrogate total reward from $\tilde{R}=\sum_{j=1}^{T}\tilde{r}_{j}$, the final importance scores can be optimized by minimizing the loss function (11).

With estimated importance scores, we can search for the optimal treatment regime under the SWL value function (9), via maximizing the objective function with a smooth convex surrogate function $\psi$ (Bartlett et al., 2006):

In particular, we can employ the logistic function as a surrogate to the indicators, i.e., $\psi(x;\lambda)=(e^{-\lambda x}+1)^{-1}$. The hyper-parameter $\lambda$ controls the growth rate where a larger $\lambda$ value makes the surrogate converge to the 0-1 indicator function faster.

To conclude, the proposed SWL method inherits the property of treatment mismatching from SAL and further enhance stage heterogeneity via the empirical stage importance scores estimated from the attention mechanism. The resulting SWL scheme is able to estimate the sequential DTR under the weighted multi-label classification framework, where optimal regime learning at stages with substantial treatment effects are prioritized. In the next section, we will present the theoretical results of the proposed SWL method.

The adoption of surrogate functions eases the optimization procedure. In this subsection, we establish the Fisher consistency between the value function $V^{SW}$ and its surrogate form $V^{SW}_{\psi}$. Specifically, we show that the optimal surrogate treatment decision $\text{sign}(f^{*}_{\psi j})$ is aligned with the optimal decision $\text{sign}(f^{*}_{j})$ at each stage. The obtained result is presented in Theorem 13.

Theorem 13 guarantees the same treatment decisions could be obtained from the surrogate estimators and the maximizing regime under the SWL scheme. This validates the usage of smooth surrogate functions to approximate the indicator functions at each decision stage which cannot be solved under the IPWE framework. In addition, according to Lemma 4 in Appendix A.3, the surrogate is only required to be an even function and produce the same sign as the treatment effect. In fact, a wide class of surrogate functions, such as indicators, logistic functions, and binary cross-entropy, satisfies such requirements and increases the optimization flexibility of the proposed SWL method.

In the following subsection, we also demonstrate the Fisher consistency between SWL and the optimal DTR $\mathscr{D}^{*}=\{\mathcal{D}_{j}^{*}\}_{j=1}^{T}$. The optimal decision $d^{*}_{t}$ on arbitrary stage $t$ from equation (1) has the form of

Compared to SWL maximizing regime $\mathbf{f}^{*}$ in equation (13), the optimal DTR $\mathscr{D}^{*}$ also aims to maximize expected future reward, but further assumes every future treatment step matches with the optimal decision. Thus, depending on the underlying behavioral distributions of rewards and treatment assignments, SWL could recover $\mathscr{D}^{*}$ asymptotically only under much more stringent conditions.

To fill in the gap between SWL and optimal DTR, we investigate the boundary condition where the SWL estimators might produce different treatment decisions from the optimal regime. This condition can be quantified via the relationships between the expected optimal reward and potential loss due to sub-optimal treatments, defined as follows:

As Theorem 2 shows, Fisher consistency depends on the dominance of the optimal treatments. For a better understanding of Condition 15, we begin with an extreme scenario where every individual receives the optimal decisions. Obviously, optimal decisions in this case are dominant as there is no other regime assigned, and Condition 15 is satisfied at every stage, i.e., $P^{\dagger}(A_{t},H_{t})=0$ and the optimal decision $d^{*}_{t}$ maximize the $R^{*}(A_{t}=d^{*}_{t},H_{t})$. Furthermore, since the expected reward $\mathbb{E}\left\{R\;\middle|\;A_{t},H_{t}\right\}$ is equal to the expected optimal reward $R^{*}(A_{t},H_{t})$, it is straightforward to show that the SWL maximizing regime $\mathbf{f}^{*}$ is the same as the optimal DTR $\mathbf{d}^{*}$ according to Equations (13) and (14).

In a more general sense, Condition 15 balances the optimal reward gain and the risk of future sub-optimal decisions when making a current-stage decision. For example, though the optimal regime could be hardly assigned to some patients, i.e., $P^{\dagger}(d^{*}_{t},H_{t})$ is large, the optimal decision $d^{*}_{t}$ is still preferred if its expected optimal reward gain, $R^{*}(d^{*}_{t},H_{t})$, is much higher and therefore dominates other regimes. On the other hand, if different decisions yield similar optimal reward gain, i.e., $R^{*}(d^{*}_{t},H_{t})-R^{*}(d^{\dagger}_{t},H_{t})$ is small, and there is a high chance of significant losses from future sub-optimal treatments after assigning the optimal treatment at the current stage, i.e., $(R^{*}(d^{*}_{t},H_{t})-R^{\dagger}(d^{*}_{t},H_{t}))\cdot P^{\dagger}(d^{*}_{t},H_{t})$ is large, the SWL estimators will choose an alternative decision which provides a reasonable future reward while minimizing the reward losses from future sub-optimal treatments. Accordingly, due to Condition 15, our proposed SWL is able to recover the optimal regime and adaptively make treatment changes based on the expected rewards, regime assignment rates, and possibilities of making sub-optimal regime from the collected empirical data.

Theorems 13 and 2 establish the consistency properties of the proposed SWL method. To investigate finite sample performance of the proposed approach, we also establish the performance error bound and investigate the convergence rate. In the following, we require measuring the function space complexity for the parameterized functional space $\mathcal{F}$.

Assumption 16 characterizes the functional space complexity with the logarithmic minimum number of balls with radius $u$ required to cover a unit ball in $\mathcal{F}$, and is satisfied under various functional spaces such as the reproducing kernel Hilbert space (RKHS) and Sobolev space (Van de Geer, 2000; Steinwart and Christmann, 2008). Consequently, the performance bound between $V^{SW}(\mathbf{f}^{*})$ and $\hat{V}^{SW}_{\psi}(\hat{\mathbf{f}}_{n})$ is provided in the following Theorem 3.

In Theorem 3, the finite-sample performance bound can be broken down into two separate bounds: the surrogate error bound between $V^{SW}$ and $V^{SW}_{\psi}$ and the empirical estimation error bound of $V^{SW}_{\psi}$. As a result, the SWL performance error convergence rate could be obtained at $\mathcal{O}(\epsilon_{n,j}+n^{-1/(2\alpha+2)})$. In particular, the first term depends on the choice of the surrogate function. When the surrogate is well-selected as a logistic function, e.g., $\psi(x;\lambda_{n})=(e^{-\lambda_{n}x}+1)^{-1}$ with a rate hyper-parameter $\lambda_{n}$, the surrogate error $\epsilon_{n,j}$ vanishes to zero at the rate of $\mathcal{O}(e^{-n})$, which is much faster than the second term; and therefore the performance error bound could be reduced to $\mathcal{O}(n^{-1/(2\alpha+2)})$.

Furthermore, based on the $\mathcal{O}(n^{-1/(2\alpha+2)})$ convergence rate, Theorem 3 provides the finite-sample upper bound which validates the estimation risk and demonstrates how SWL converges under different parametric space settings. For instance, when the historical information $\mathcal{H}$ is an open Euclidean ball in $\mathbf{R}^{d}$ and the functional space is specified as the Sobolev space $\mathbb{W}^{k}(\mathcal{H})$ where $k>d/2$, one can choose $\alpha=d/2k$ to obtain an error upper bound of rate at $\mathcal{O}(n^{-d/2(d+2k)})$. In addition, when the functional space is finite, the upper error bound could reach the best rate at $\mathcal{O}(n^{-1/2})$ as $\alpha\rightarrow 0$, which achieves the optimal rate provided in the literature (Zhao et al., 2015, 2019).

To summarize, our finite-sample performance error bound recovers the best-performing convergence rate found in the existing literature and meanwhile provides a non-asymptotic explanation of the proposed multi-stage DTR method in empirical settings. With different choices of metric entropy, the performance error bound can be flexibly adapted to various functional spaces and is not limited to the RKHS discussed in Zhao et al. (2015). In addition, our metric condition from Assumption 16 only needs to be satisfied under a more relaxed empirical $L_{2}$-norm on the collected samples $H_{1:n}$, compared to the supremum norm assumption in Zhao et al. (2019).

Sample sizes and the number of decision stages are the two important factors that affect the curse of full matching as introduced in section 2.3. To fully examine the effects of these two factors on our proposed algorithm, we first conduct simulations on sample sizes $n=500,1000,5000$, and the number of decision stages $T=5,8,10$ where treatments $A_{t}$ are randomly matched with the nonlinear optimal decisions at $50\%$ chance and the number of important stages is set to $0$. The results of model performance are summarized in Table 2.

According to Table 2, we notice that SAL/SWL outperform all other competing methods with respect to the estimated total rewards. Given a fixed number of stages, every single model deteriorates as expected when the sample size decreases as indicated by the smaller values of estimated total rewards, but the improvement margin of the proposed methods compared to the best-performing competing method increases. In particular, when $T=5$, the SAL/SWL improves the estimated total rewards nearly three times as much, from $10.08\%$ to $27.92\%$, as the sample size decreases from $5000$ to $500$. In addition, the difference between our model performance and other competing methods enlarges with an increasing number of decision stages. For instance, when $n=5000$, the improvement rates increase from $10.08\%$ to $20.54\%$ as $T$ grows from $5$ to $10$. This implies that the proposed method has a more efficient utilization of the observed information and more advantages when the sample size is small and the number of treatment stages is large.

We illustrate the curse of full matching under various sample sizes and numbers of stages. However, the empirical dilemma could compromise the convergence of the DTR methods at the same time. To describe a more straightforward association between full-matching rates and model performance while minimizing the effects from non-convergent results, we set $n=5000$, $T=10$, and directly adjust the matching probabilities between the assigned and optimal treatments. Specifically, at each stage, we consider the assigned treatment with a 50%, 60%, 70%, and 80% probability of matching the optimal treatment. As a result, the number of stage receiving optimal decisions follows a binomial distribution with these respective probabilities, and the full-matching rates are $0.5^{10}\approx 0.001$, $0.6^{10}\approx 0.006$, $0.7^{10}\approx 0.03$, and $0.8^{10}\approx 0.1$.

Based on the results presented in Table 3, we observe that SAL/SWL outperform the rest of the competing methods, with greater improvement over BOWL as the full-matching probability decreases. This is expected since BOWL, as an IPWE-based method, depends heavily on full-matching treatments for convergence, while our methods incorporate sub-optimal treatments and thus are more robust across varying treatment matching scenarios. Among competing methods, the Q-learning performs poorly due to its dependence on a correctly specified outcome model, which is challenging with high heterogeneity and a large number of stages. AIPWE and RWL face similar challenges with their augmented unbiased and residual terms, and AIPWE further suffers from its complex estimator format. The C-learning, which combines regression-based methods with AIPWE policy search techniques, shows the highest performance among the competing methods when the full-matching rates are small. Nonetheless, both the regression and robust policy-search components increase computational complexity and rely on outcome models, parametric or non-parametric, to stabilize the IPWEs. In contrast, the proposed methods provide a fundamental solution to break the curse of full-matching, effectively enhance sample efficiency, and achieve superior model performance.

In this numerical study, we are interested in analyzing the sensitivity of our proposed method to the functional complexity of the underlying optimal regime. Specifically, we also consider the linear treatment regimes and include the homogeneous decision rule setting where the optimal rules are the same at all time points, i.e. $f_{j}^{*}=f_{1}^{*}$ for all $2\leq j\leq T$. Note that the homogeneous rules can be oftentimes encountered in a high-frequency treatment session as the optimal regime is unlikely to update in a short period of time. We combine the results of four settings when $T=5$ in Figure 2.

As algorithms converge, a decrease in the performance improvement rate/slope is expected. According to this criterion, we observe that the proposed SAL/SWL methods have similar convergence rates and converge to Oracle faster than all other methods. For instance, under the linear homogeneous treatment rule setting, SWL achieves an averaged $97.36\%$ matching accuracy, compared to $74.71\%$ for BOWL when the sample size reaches $1000$. Though the presented empirical results can be affected by the implementation and choices of hyper-parameters, based on the similar performance increasing rate between SWL and BOWL when the sample size is larger than $2000$, we can confirm with our theoretical results that our SWL reaches the same state-of-the-art asymptotic convergence rate as BOWL. In addition, from the deteriorated performance in the nonlinear heterogeneous setting compared to the linear homogeneous setting, we verify that the increasing level of functional complexity raises the optimization difficulties and hinders the model convergence rate with limited empirical examples. Nevertheless, our proposed method outperforms all competing methods by a considerable margin, especially when the sample size is small.

In this subsection, we illustrate the advantages of incorporating stage heterogeneity with the stage importance scores. We adjust the level of heterogeneity by changing the number of important stages, where fewer important stages induce stronger heterogeneity. In addition, to maximize the heterogeneity among stages, we consider the linear homogeneous decision rule setting as illustrated in the previous simulation subsection. Figure 3 provides the obtained results when $n=500$ and $T=10$.

As shown, while the proposed SAL still outperforms the rest of the competing methods under the first three scenarios, the importance scores can further improve the performance of SAL with a greater margin when the heterogeneity among the stages gets stronger. Moreover, the stability of the SAL can be improved with the stage weights when strong stage heterogeneity exists. We conclude that the proposed stage importance scores are able to explicitly incorporate stage heterogeneity into the SAL estimator and can be used for regime searching on stages which contribute to improving treatment effects.