Post Reinforcement Learning Inference

Post Reinforcement Learning Inference

\RUNTITLE{Post Reinforcement Learning Inference}

\TITLE{Post Reinforcement Learning Inference}

\ARTICLEAUTHORS{ \AUTHOR{Vasilis Syrgkanis\quad Ruohan Zhan}

\RUNAUTHOR{Syrgkanis and Zhan} }

\ABSTRACT{ We study estimation and inference using data collected by reinforcement learning (RL) algorithms. These algorithms adaptively experiment by interacting with individual units over multiple stages, updating their strategies based on past outcomes. Our goal is to evaluate a counterfactual policy after data collection and estimate structural parameters, such as dynamic treatment effects, that support credit assignment and quantify the impact of early actions on final outcomes. These parameters can often be defined as solutions to moment equations, motivating moment-based estimation methods developed for static data. In RL settings, however, data are often collected adaptively under nonstationary behavior policies. As a result, standard estimators fail to achieve asymptotic normality due to time-varying variance. We propose a weighted generalized method of moments (GMM) approach that uses adaptive weights to stabilize this variance. We characterize weighting schemes that ensure consistency and asymptotic normality of the weighted GMM estimators, enabling valid hypothesis testing and uniform confidence region construction. Key applications include dynamic treatment effect estimation and dynamic off-policy evaluation. }

\KEYWORDS{reinforcement learning, GMM estimators, adaptive weighting, asymptotic normality, strong Gaussian approximation, hypothesis testing, dynamic treatment effects, dynamic off-policy evaluation }

Introduction

Adaptive data collection has become a staple of the digital economy. Most major digital platforms invoke adaptive experimentation algorithms to optimize their service. Adaptive experimentation allows one to progressively update their experimentation strategy and leads to efficient sample usage chu2011contextual,agrawal2013thompson and efficient use of the experimentation budget for increased hypothesis testing power russo2016simple. Moreover, frequent adaptive experimentation is starting to become adopted in other domains such as personalized healthcare murphy2005experimental,offer2021adaptive. The popularity of adaptive experiments has increased the availability of data collected from such designs. However, adaptive data collection raises many new research challenges, especially in the case of post-collection statistical analysis. For instance, constructing confidence intervals and testing statistical significance for alternative candidate policies or for structural parameters like average treatment effects, from adaptively collected data, has been shown to be theoretically challenging.

Prior work has mostly focused on the bandit setup, where at each time the experimenter only interacts with units sampled from the environment once and observes the immediate outcome deshpande2018accurate,zhang2021statistical,hadad2021confidence,bibaut2021post,zhan2021off. This unfortunately cannot accommodate many applications including adaptive clinical trials and dynamic treatment regimes, where a patient often receives multiple rounds of treatments to improve the outcome murphy2003optimal,lei2012smart, or a digital platform interacts with their users over a sequence of multiple page visits chen2019top.

The goal of our work is to address this largely un-explored area of post-adaptive data collection inference, from such adaptive experiments that, within each experiment phase, involve multiple interactions with the same treated unit. In particular, we consider data that are collected from reinforcement learning (RL) algorithms sutton1998introduction and provide estimation and inference for structural parameters of interest under semi-parametric assumptions neyman1979c,laan2003unified,chernozhukov2022locally.

In RL, an adaptive experimentation algorithm (from now on “the agent”) interacts with units that are sampled independently and identically (i.i.d.) from the environment. For each unit, the agent observes an initial state, then applies multiple treatments that cause state transitions, and finally observes an outcome at the end.\footnote{We focus on cases where only the final outcome is observed. Our framework can be generalized to settings where intermediate outcomes are also revealed.} We term the sequence of interactions with a unit as an “episode”. For example, consider an educational platform aiming to encourage users to enroll in a course. The platform, acting as the agent, might first interact with a user on the homepage, then guide them to an enrollment page, and finally to a payment page, with each step representing a state transition. Throughout these stages, the platform can experiment with different page designs (treatments) to achieve the final goal of course enrollment.

The agent's behavior policy is typically adapted over time to incorporate learning from past realizations. In particular, we focus on a broad class of adaptive data collection mechanisms, where the agent interacts with a batch of units before updating its policy for subsequently arriving units. This setting encompasses the episodic reinforcement learning framework neu2020unifying and generalizes the classical bandit framework, in which the agent interacts with each unit for only a single round in an adaptive manner.

Such adaptivity progressively improves the agent's performance but results in a nonstationary behavior policy and introduces dependence between observations. As a result, we cannot simply view the data from each episode as i.i.d. samples and pass them to traditional estimation pipelines lewis2020double. Even when employing estimation techniques designed to address time-series correlation, these often require stationarity—a condition not met by most adaptive experimentation processes. This nonstationarity causes evolving discrepancy between the behavior and target policies, an issue known as changing “overlap” and resulting in varied estimation variances for sequentially collected samples imbens2004nonparametric,hadad2021confidence. Therefore, averaging samples uniformly can be suboptimal, resulting in significant variance and a non-normal asymptotic distribution, which complicates post-experimental inference. This problem, evident even in single-interaction bandit scenarios, has prompted recent literature to suggest re-weighting samples to stabilize time-varying variance, such that the resulting estimators are consistent and asymptotically normal deshpande2018accurate,hadad2021confidence,zhang2021statistical,zhan2021off,bibaut2021post.

RL settings further complicate the estimation problem in two ways. First, exogenous random shocks during state transitions within an episode affect subsequent states and are correlated with future treatments; these shocks, being unobservable, cause unmeasured confounding robins1986new,robins2004optimal,chakraborty2013semi. To address the identification issue, we follow the semi-parametric inference literature, formulating the dynamic treatment effect estimation problem as estimating the structural parameters in a structural nested mean model robins2004optimal,lok2012impact,vansteelandt2016revisiting. In our model, each stage is linked to a specific structural parameter that quantifies the treatment effect at that stage. We demonstrate that these parameters are the solutions to stage-wise moment equations, derived through $g$-estimation and constructed in reverse order, from the last stage to the first. Given the nonstationarity of RL data, which often lead to time-varying estimates across units as discussed above, we apply nonuniform and adaptive weights to the unit samples to stabilize these variances.

Second, estimation with RL data often faces the “curse of horizon” challenge, where the overlap between behavior and target policies deteriorates across episodic stages, leading to accumulated estimation variance. Prior work suggests heuristics like weight clipping, which---while effective in controlling variance---introduce a small bias precup2000eligibility,chen2019top. However, our approach, through nuanced modeling, ensures that moment equations across stages remain uncorrelated with zero covariance. This allows for the application of stage-wise adaptive weights, which stabilize the variance at each stage based on information available up to that point, effectively circumventing the problem of variance accumulation.

This work also enriches the inference literature when using adaptively collected data by providing results for GMM-estimation. Prior work mostly addresses inference in the context of M-estimation deshpande2018accurate,zhang2021statistical, i.e. parameters that can be defined as the minimizers of a population loss function. However, structural parameters in problems defined by moment conditions, which often arise in the dynamic treatment regime settings, cannot be phrased as M-estimators (rather they can be better thought as instrumental variable problems).

Our Contributions

Our main contributions are outlined as follows. First, we propose a weighted GMM-estimator using adaptive RL data to estimate structural parameters of counterfactual policy values and dynamic treatment effects. These parameters are defined by moment conditions within a dynamic treatment regime, modeled through a semiparametric structure nested mean model for each unit's episodic data. Specifically, for a target counterfactual policy $\pi^*$, we aim to estimate its policy value $\theta_0^*$---the expected final outcome $Y$ under policy $\pi$. We further define $\theta_j^*$ as the stage-wise dynamic treatment effect of treatment $T_j$ at stage $j$, reflecting the expected change in the final outcome $Y$ due to $T_j$, assuming future treatments follow $\pi^*$. Let $L$ be total number of stages in an episode for each unit. Our objective is to derive the structural parameter set $\theta^*=(\theta_0^*, \theta_1^*,\dots, \theta_L^*)$ from RL data, capturing both the policy value and stage-wise dynamic treatment effects.

We hereby illustrate the core concept of our approach under the episodic RL setup. We consider $n$ sequentially collected episodes $\{Z_1, \dots, Z_n\}$. Let $\mathcal{F}_i:=\sigma(Z_{1:i-1})$ represent the $\sigma$-algebra with respect to which the random variables $Z_{1:i-1}$ are measurable. The true parameter vector $\theta^*$ satisfies the moment condition $\mathbb{E}[m(Z_i;\theta^*) \mid \mathcal{F}_i] = 0$, where $m(\cdot)$ denotes bounded linear moment functions derived from our semiparametric model. To enhance the stability and efficiency of the estimation, we adopt the generalized method of moments (GMM) throughout this paper.\footnote{The results presented here can be directly extended to Z-estimation.} The Jacobian of the moment function, denoted by $J_i := \nabla_{\theta} m(Z_i; \theta^*)$, is nonstationary across episodes due to adaptivity. As a result, the traditional (unweighted) GMM estimator, which is based on minimizing the norm of empirical moment equations, may exhibit unstable asymptotic behavior. To address this, we introduce adaptive weights $H_i$, measurable with respect to $\mathcal{F}_i$, to stabilize the estimation.

Let $\hat{\theta}_n$ denote the weighted GMM estimator that minimizes $\left\|\frac{1}{n}\sum_{i=1}^n H_i m (Z_i;\theta)\right\|_A^2$, where $A$ is a positive-definite matrix and $\|\cdot\|_A$ is the norm induced by $A$. When $\hat{\theta}_n$ lies in the interior of $\Theta$, where $\Theta$ is the bounded parameter space of $\theta^*$, it satisfies the first-order condition: $\{\frac{1}{n}\sum_{i=1}^n H_iJ_i\}A\{\frac{1}{n}\sum_{i=1}^n H_im (Z_i;\hat\theta_n)\}=\mathbf{0}$. By Taylor expansion, we have:

The property of $\hat{\theta}_n$ thus depends on the behavior of $B_n$ and $S_n$, both of which can be stabilized by the choice of weights $\{H_i\}$.

Our second contribution is to identify generic weighting schemes to achieve consistency and asymptotic normality of the weighted GMM-estimator $\hat{\theta}_n$. The choice of weights $\{H_i\}$ should address two critical aspects: ensuring the normalizing matrix $B_n^\top A B_n$ in (ref) remains well-posed to prevent the explosion of the estimated parameter $\hat{\theta}_n$, and the right-hand side $B_n^\top A S_n$ converges to a Gaussian distribution for asymptotic normality. We show that $S_n$ is the sum of a martingale difference sequence, and we establish its convergence by leveraging recent advances in martingale limiting theorems under a crucial homoscedasticity assumption. We note that simply providing a central limit theorem for (ref) is not enough to achieve inference results on policy value $\theta_0^*$ or dynamic treatment effect $\theta_j^*$, since the normalizing matrix $B_n^\top A B_n$ is likely to not concentrate by the sheer nature of adaptivity. Addressing this, we adapt recent work in cattaneo2022yurinskii, which develops strong Gaussian approximation results for martingale data, to our context and show that (ref) converges to a Gaussian distribution at a uniform rate. This allows us to invert the normalizing matrix $B_n^\top A B_n$ and approximate the estimation error, $\hat{\theta}_n-\theta^*$, with a Gaussian distribution, paying the way for inference results that are the primary focus of this paper.

Our third contribution is to substantiate the generic weighting schemes and offer specific weighting choices to achieve consistency and asymptotic normality for a wide range of RL algorithms with polynomially decaying exploration rates. We focus on bilinear parametrization, which includes common setups such as categorical treatment and polynomial scalar treatment. We show that the adaptive weights for consistency can be directly derived from the collected data. Further, we identify the “oracle weights” to achieve strong Gaussian approximation. While these weights may require additional estimations, the necessary quantities to be estimated only depend on the state transition dynamics, which are independent of the behavior policy; this allows for their consistent estimation through online regression techniques, which we term “feasible weights”. We prove that the GMM-estimator, when applied with these feasible weights, preserves strong Gaussian approximation, thereby facilitating post-RL inference.

Finally, we apply our estimation and inference framework to high-dimensional Markovian models. We offer estimation guarantees for feasible weights and show that the adaptively weighted GMM-estimator, when using these feasible weights, achieves strong Gaussian approximation. This approximation enjoys a uniform convergence rate of $O(n^{-(1-L\alpha)/10})$, where $L$ represents the episode length, and $\alpha$ indicates the exploration decaying rate of the RL agent. We further substantiate our theoretical guarantees with numerical evidence. Our findings show that the standard GMM-estimator, when unweighted, deviates from asymptotic normality, resulting in either insufficient coverage or overly conservative inferences for structural parameters. Conversely, our weighted GMM-estimator, applied with feasible weights, consistently achieves near-perfect coverage, tighter confidence intervals, and more precise estimates for evaluation policy values. Its performance closely matches that of estimations under oracle weights, which, though ideal, require inaccessible true structural parameters and are generally impractical. Furthermore, we demonstrate our method's robustness under model misspecification, employing a polynomial model to approximate the true exponential model within our semiparametric framework. Even under misspecification, our weighted GMM-estimator with feasible weights outperforms the standard approach, offering more accurate policy value estimates and maintaining near-nominal coverage as approximation complexity increases. This robustness highlights our method's suitability for real-world scenarios, where model misspecification is often inevitable.

Setup

This section establishes our problem framework by detailing the data generating process. We start by describing the stochastic control process, which involves multiple stages of interaction with individual units. Subsequently, we define the structural parameter estimation problem within this framework.

Episodic Potential Outcome

We start with describing the stochastic control process to roll out an episode for a unit and the corresponding potential outcome. Consider each unit starts from an initial state $S_1$ i.i.d. sampled from a fixed distribution $P_S$. An RL agent then runs an episode of $L$ stages for each unit. At each stage $j \in [1:L]$, the agent observes the unit’s current state $S_j \in \mathcal{S} \subset \mathbb{R}^{d_s}$ and then assigns a treatment $T_j \in \mathcal{T} \subset \mathbb{R}^{d_\tau}$ (also known as an action or intervention), where $d_s$ and $d_\tau$ denote the dimensions of the state and treatment, respectively. The treatment transitions the unit to a new state $S_{j+1}$ for the next stage. The long-term outcome of interest, denoted by $\ensuremath{Y} \in \mathcal{Y} \subset \mathbb{R}$, is observed only at the end of the episode $(S_1, T_1, S_2, \dots, S_L, T_L)$.

Following the potential outcome framework, we denote the potential outcome under a sequence of treatment assignments $\{\tau_1,\dots, \tau_L: \tau_j\in\mathcal{T}\}$ as $\ensuremath{Y}(\tau_{1:L})$.\footnote{We use $A_{n_1:n_2}$ as a short hand for the set $\{A_{n_1},\dots, A_{n_2}\}$ for any indexed variable $A_i$ and any two integers $n_1\leq n_2$.} A dynamic treatment assignment policy $\pi = \{\pi_1, \dots, \pi_L\}$ specifies a sequence of treatment rules applied across stages of an episode. Specifically, each $\pi_j$ determines the treatment $T_j$ at stage $j$ based on the history $(S_{1:j}, T_{1:j-1})$, with $\pi_j : \mathcal{S}^j \times \mathcal{T}^{j-1} \rightarrow \mathcal{T}$. We overload the notation and use $\ensuremath{Y}(\pi_{1:L})$ to denote the potential outcome under policy $\pi$. We introduce an assumption critical to our analysis: all confounders influencing both the treatment assignment and the final outcome are observed by the states.

This assumption is an extension of the unconfoundedness assumption in causal inference literature to settings where a unit undergoes multiple stages of treatment rosenbaum1983central,robins2004optimal.

Inference Goal: Value of Dynamic Treatment Assignment Policy

Given a policy $\pi^*$ to be evaluated, the goal in this paper is to estimate its policy value, which is defined as the mean outcome under this policy, outlined below:

Inference on $\theta^*_0$ allows one to argue whether the observed outcome achieved under the behavior policy is statistically larger than the outcome under some simple baseline policies, and therefore whether the data collection agent (for example, from some RL algorithm) leads to any statistically significant benefits as compared to these baseline policies.

We introduce the following definition to attribute the long-term outcome to intermediate treatments at each stage.

In other words, the blip effect $\gamma_j^{(\pi^*)}(\cdot)$, given the realized history $(S_{1:j}, T_{1:j-1})$ and conditional on assigned treatment $T_{i,j}$, captures the effect of assigning treatment $T_{i,j}$ in the current period relative to a baseline control policy $0 \in \mathcal{T}$ that applies the status quo treatment, assuming future treatments follow policy $\pi^*$. From now on, we fix the evaluation policy $\pi^*$ and omit the superscript in the blip function, i.e. we use $\gamma_j(\cdot)$ to denote $\gamma_j^{(\pi^*)}(\cdot)$ for notation convenience.

To identify the policy value of $\pi^*$ from the observed data, we introduce a semi-parametric structural assumption on the blip functions, following lewis2020double,robins2004optimal.

Note that even when the blip functions are assumed to be parametrically linear, we do not make any assumption on the conditional mean of the evaluation policy outcome $\mathbb{E}\left[\ensuremath{Y}(\pi^*_{1:L})\mid S_1\right]$, which can be potentially nonlinear in the initial state $S_1$. In later sections we shall use moment equations to identify the unknown structural parameter $\theta_j$, and this assumption implies that the moment conditions are linear in target parameters. This assumption allows for a wide range of model classes including high-dimensional Markovian models substantiated in Section (ref).

To this end, we can explicitly characterize the policy value $\theta^*_0$ for evaluation policy $\pi^*$, following results in robins2004optimal -- which we adapt to our setup and provide its proof (along with all other proofs) in the appendix for completeness.\footnote{This result can be viewed as an analogue of the performance difference lemma that is frequently used in the reinforcement learning literature kakade2002.}

Research Question

We conclude this section by formally describing our problem within the context formalized above.

\paragraph{Adaptive Episodic RL Data.} We consider the setting where an RL agent sequentially rolls out episodes for $n$ units, each over $L$ stages. Before episode $i$ begins, the agent updates the treatment assignment policy $\pi_i := \{\pi_{i,1}, \dots, \pi_{i,L}\}$ based on data from previous episodes. Let $Z_i := \{S_{i,1}, T_{i,1}, \dots, S_{i,L}, T_{i,L}, Y_i\}$ denote the data collected during episode $i$. The distribution of $Z_i$ generally depends on the past realizations $Z_{1:i-1}$. The full dataset consists of observations $\{Z_i\}_{i=1}^n$. Our results directly extend to batched episodic RL settings, where the RL agent may interleave interactions with units within a batch and simultaneously progress these units through their respective episodes. This extension is possible because our technical analysis is based on filtrations defined with respect to unit-level episodes. As long as the behavior policy remains fixed during the rollout of each individual unit, similar filtrations can be constructed based on episode completion, and our results continue to hold.

We assume that the behavior policy $\pi_i$ used to generate data for episode $i$ is known and recorded by the RL agent. Without this information, recovering the behavior policy from the data is generally difficult or even statistically impossible, since each $\pi_i$ may correspond to only one observed episode. This assumption is also common in prior work on inference with adaptively collected bandit data hadad2021confidence,zhang2021statistical.

\paragraph{Goal.} Given the data, we aim to estimate and perform inference on the structural parameters in the parameter space $\Theta\subset \mathbb{R}^{1+dL}$: \[\theta^* = (\theta_0^*, \theta_1^* ,\dots,\theta_L^*)\in\Theta,\] where $\theta_0^*$ corresponds to the value of the evaluation policy $\pi^*$, and $\{\theta_j^* \in \mathbb{R}^d:j\in[1:L]\}$ denote the dynamic treatment effect parameters that appear in the linearization of the blip functions. In particular, we assume that $\theta^*$ is an interior point of $\Theta$.

\paragraph{Notation.} We use $i$ to index episode and $j$ to index stage within an episode. Define the $\sigma$-fields $\mathcal{F}_{i} := \sigma(\{Z_{1:i-1}\})$ and $\mathcal{F}_{i,j} := \sigma(\{Z_{1:i-1}, S_{i, 1:j-1}, T_{i,1:j-1}\})$ with the convention that $\mathcal{F}_{i,1}\equiv \mathcal{F}_{i}$ and $\mathcal{F}_{i,0}\equiv \mathcal{F}_{i}$. We also introduce an augmented $\sigma$-field: $\mathcal{F}_{i,j}^+$, which extends $\mathcal{F}_{i,j}$ by including the $\sigma$-algebra generated by the next state $S_{i,j}$; this additional information about state $S_{i,j}$ is available prior to the treatment assignment by the agent for unit $i$ at stage $j$. We use $\mathbb{E}_{i}[\cdot]$ as a short hand for the conditional expectation $\mathbb{E}[\cdot\mid \mathcal{F}_{i}]$ and use $\mathbb{E}_{i,j}[\cdot]$ for $\mathbb{E}[\cdot \mid \mathcal{F}_{i,j}]$, $\mathbb{E}_{i,j}^+[\cdot]$ for $\mathbb{E}[\cdot \mid \mathcal{F}_{i,j}^+]$. Similarly, we define $\mathrm{Var}_{i}(\cdot), \mathrm{Var}_{i,j}(\cdot)$ and $\mathrm{Var}_{i,j}^+(\cdot)$ as shorthand for the conditional variances $\mathrm{Var}(\cdot\mid \mathcal{F}_{i}), \mathrm{Var}(\cdot\mid \mathcal{F}_{i,j})$ and $\mathrm{Var}(\cdot\mid \mathcal{F}_{i,j}^+)$, respectively. For a random vector $X$, we define $\mathrm{Var}_i(X)=\mathbb{E}_i[XX^\top] - \mathbb{E}_i[X]\mathbb{E}_i[X]^\top$, $\mathrm{Var}_{i,j}(X)=\mathbb{E}_{i,j}[XX^\top] - \mathbb{E}_{i,j}[X]\mathbb{E}_{i,j}[X]^\top$ and $\mathrm{Var}_{i,j}^+(X) = \mathbb{E}_{i,j}^+[XX^\top]-\mathbb{E}_{i,j}^+[X] \mathbb{E}_{i,j}^+[X]^\top$.

Identification and Estimation

In this section, we provide conditional moment equations for identifying the target parameters. We then propose our adaptively weighted method of moments (AW-GMM) for parameter estimation using the adaptively collected RL data.

Identification via Moment Equations

We construct moment equations to identify $\theta^*$, as formalized by the lemma below. This approach is known as the $g$-estimation framework, typically used in the context of structural nested mean models robins2004optimal. However, most prior work focuses on data collected by a fixed behavior policy robins2004optimal,chakraborty2013semi,lewis2020double, while here we apply the $g$-estimation approach to adaptive RL data and show that the true parameter vector $\theta^*$ also satisfies the moment restrictions proposed in robins2004optimal.

The moment conditions in the above lemma can be expressed more compactly as a single vector of moment constraints. Given a unit $i$ with episodic data $Z_i$, define $\beta_i:=(\beta_{i,0},\beta_{i,1},\dots, \beta_{i,L})$, where each component is given by $\beta_{i,j}=Y_i(\Psi_{i,j} -\bar \Psi_{i,j})$; define

Then for any parameter vector $\theta$, the moment function $m(Z_i;\theta)$ can be written as:

In particular, the moment condition in (ref) can be constructed directly from observed data. The key quantity, $\bar\Psi_{i,j} = \mathbb{E}_{i,j}^+[\Psi_{i,j}]$, is the conditional expectation over the treatment assignment $T_{i,j}$ under the behavior policy. Since the behavior policy is assumed to be known, this expectation can be computed exactly. As a result, the moment function $m(Z_i; \theta)$ is also computable from the observed data.

Finally, Lemma (ref) implies that the true model parameter vector $\theta^*$ solves the expected moment equation:

where $\bar\beta_i:=(\bar\beta_{i,0},\bar\beta_{i,1},\dots, \bar\beta_{i,L})$ for $\bar\beta_{i,j}=\mathbb{E}_{i,j}^+\left[Y_i(\Psi_{i,j} -\bar \Psi_{i,j})\right]$, and

This observation will become handy in our subsequent analysis.

Adaptively Weighted Generalized Method of Moments (AW-GMM)

Given the moment conditions in (ref), one can estimate the target parameter $\theta^*$ using moment-based methods such as Z-estimation or the generalized method of moments (GMM). To improve stability and efficiency (particularly when $\Psi$ has higher dimension than $\Phi$), we focus on GMM throughout this paper. The results also apply directly to Z-estimation. Specifically, the standard GMM estimator $\tilde{\theta}_n$ minimizes the norm of the empirical moment equations:

where $A$ is a positive-definite matrix and $\|\cdot\|_A$ denotes the norm induced by $A$. When $\tilde{\theta}_n$ lies in the interior of $\Theta$, it satisfies the first-order condition: $\{\frac{1}{n}\sum_{i=1}^n J_i\}A\{\frac{1}{n}\sum_{i=1}^n m (Z_i;\tilde\theta_n)\}=\mathbf{0}$.

To build more intuition, consider the case where the sample average Jacobian $\frac{1}{n}\sum_{i=1}^n J_i$ is invertible. Then the first-order condition implies that the GMM estimator solves the empirical moment equations such that: $\frac{1}{n}\sum_{i=1}^n m (Z_i;\tilde\theta_n)=\mathbf{0}$. The asymptotic behavior of $\tilde{\theta}_n$ is typically analyzed through a Taylor expansion around the true parameter $\theta^*$, which links the normalized estimation error $J_i(\tilde{\theta}_n-\theta^*)$ with an empirical influence function $m(Z_i; \theta^*)$:

However, since the behavior policy is adaptively evolving over time, the Jacobian $J_i$ is likely to diverge, and the variance of the empirical influence functions on the right-hand-side might fluctuate and never converge. This risk of instability motivates us to modify the GMM estimator by applying non-uniform and time-varying adaptive weights, which are carefully chosen to stabilize the variance of the empirical influence functions $m(Z_i;\theta^*)$ over time, thereby preserving the desirable asymptotic characteristics of the GMM-estimator.

We hereby propose our adaptively weighted generalized method of moments (AW-GMM) estimator $\hat{\theta}_n$, which minimizes the norm of a non-uniform average of empirical moment equations:

where $H_{i}$ denotes the time-varying weighting matrix designed to counterbalance the potential divergence in the empirical influence functions. Specifically, we consider a block-diagonal structure for $H_{i}$: $H_i:= \mbox{diag}\left\{H_{i,0}, H_{i,1}, \dots, H_{i,L}\right\}$, where $H_{i,0}\in\mathbb{R}$ and $\{H_{i,j}\in \mathbb{R}^{d'\times d'}$ for each $j\in [1:L]\}$; each diagonal block $H_{i,j}$ is adapted to filtration $\mathcal{F}_{i,j}^+$ and stabilizes the empirical influence functions of data $Z_i$ involved in estimating the structural parameter $\theta_j^*$. Similar weighting schemes have been introduced in the bandit setups deshpande2018accurate,hadad2021confidence,zhang2021statistical,bibaut2021post,zhan2021off, though none applies to RL data with more than one stages.

The technical challenge in identifying proper weights for RL data is two-fold. First, it is crucial that these weights remain stable and do not become degenerate or explode, otherwise the resulting estimator would suffer from diverging variance. This is known as the “curse of horizon” challenge, which arises from the deteriorating overlap between the behavior and evaluation policies over lengthy episodes. Previous RL literature on weighting-based offline evaluation has often turned to heuristic solutions, such as weight clipping, to manage variance, albeit at the expense of introducing bias precup2000eligibility,chen2019top. We show that by carefully designing our weighting matrix $H_i$, our estimator not only controls estimation variance but also achieves asymptotic unbiasedness. Second, we want to construct these weights using information only up to current observations, following the practice of constructing weights for adaptive bandit data hadad2021confidence. By doing so, we can employ martingale limiting theorems to show that the weighted GMM-estimator admits asymptotic normality and achieves post-RL inference.

Consistency of AW-GMM Estimation

We begin by introducing a general weighting scheme that ensures consistent estimation of target parameters from adaptively collected RL data. We then apply this framework to specific RL algorithms later in the section. Recall that the weighted GMM estimator $\hat{\theta}_n$ minimizes the empirical loss given by:

By the conditional moment equation (ref) and weight $H_{i,j}$ adapting to $\mathcal{F}_{i,j}^+$, we have the true $\theta^*$ satisfies $H_i(\bar\beta_i+\bar J_i\theta^*)=\mathbf{0}$ and thus is a minimizer of the norm of averaged conditional moments:

with $\mathcal{L}(\theta^*)=0$. To establish that $\hat{\theta}_n$ is consistent for $\theta^*$, it suffices to show two conditions: (i) for every $\theta\in\Theta$, the GMM loss $ \mathcal{L}_n(\theta)$ converges to the expected loss $ \mathcal{L}(\theta)$; and (ii) the loss $\mathcal{L}(\theta)$ has a unique minimizer. The following weighting property ensures both conditions are satisfied.

In particular, Properties (ref)(a) & (ref)(b) ensure that $ \mathcal{L}(\theta)$ has a unique minimizer with high probability, and Property (ref)(c) ensures that $\mathcal{L}_n(\theta)$ uniformly converges to $ \mathcal{L}(\theta)$ for any bounded $\theta\in\Theta$.

Consider that the data is collected by an RL agent with a fixed amount of exploration, such as the $\epsilon$-greedy RL algorithms with a constant $\epsilon$. Then the uniform weighting with $H_i$ being the identity matrix satisfies Property (ref) with $\alpha_1=\alpha_2=0$, and Theorem (ref) recovers the $n^{-1/2}$ convergence rate, consistent with results in lewis2020double for i.i.d. data. However, our result is stronger in the sense that we allow for adaptivity in the experiments, as opposed to the i.i.d. settings in lewis2020double.

Consistency Weights for Decaying Exploration under Bilinear Features

We now provide explicit weighting schemes for common RL algorithms with polynomially decaying exploration rates. We particularly focus on scenarios where the instrument $\psi$ adopts the feature map $\phi$ such that $\Psi_{i,j}=\Phi_{i,j}$, and the feature map $\phi$ can be expressed as a bilinear form of $\phi_j(S_{i,1:j}, T_{i,1:j}) = \chi_{j}(S_{i,1:j}, T_{i,1:j-1}) \otimes \mu(T_{i,j})$, with $\otimes$ being the Kronecker product.\footnote{We remind that for two vectors $v\in R^p$ and $u\in R^d$, the Kronecker product is the vector whose entries contain the product of all pairs of entries of the two vectors, i.e., $v\otimes u = \text{vec}( uv^\top) = (v_1 u_1, \ldots, v_1 u_d, \ldots, v_p u_1, \ldots, v_p u_d)$.} Such feature maps accommodate a wide family of categorical treatments and continuous treatments (we defer the examples to the end of this section). Without loss of generality, we rename $\mu(T_{i,j})$, the transformation of treatment, as $T_{i,j}$ and focus on the below form throughout:

Let $d_\tau$ be the dimension of treatment $T_{i,j}$, and we use $X_{i,j}$ as a shorthand for $\chi_j(S_{i,1:j}, T_{i,1:j-1})$. Thus we can also write that:\footnote{Here we leverage the property that for two vectors $u\in \mathbb{R}^p$ and $v\in \mathbb{R}^d$, $v\otimes u=(I_d \otimes u) \cdot v$, where $I_d \otimes u = \text{diagonal}(u, \ldots, u)$. Hence, our feature map can equivalently be written as $\Phi_{i,j} = (I_{d_{\tau}} \otimes \Psi_{i,j}) \cdot T_{i,j}$.}

To identify the structure parameter $\theta^*$, one should expect sufficient overlap condition in the behavior policy and the state transition dynamics, due to co-linearities in the linear system that identifies the structural parameters. The following assumption formalizes the overlap condition.

Assumption (ref)(a) imposes an explicit rate on the decay of the amount of randomization (exploration), which is employed by the behavior policy of the RL agent as a function of the number of observations it has collected so far. Similar assumption is also required in ex post inference when using bandit data hadad2021confidence,zhan2021off. In particular, when $\alpha=0$, this becomes a relaxed version (by allowing for the dependence among observations) of the commonly made “overlap” condition in the causal inference literature on i.i.d. samples imbens2004nonparametric. Assumption (ref)(b) says that the state transition dynamic, which is independent of the behavior policy, is non-degenerate. Specifically, conditional on the past states and actions $(S_{i,1:j-1}, T_{i,1:j-1})$, the next state feature $X_{i,j}$ is not collinear; it can be achieved if the exogenous noise injected at each stage has a full-rank covariance matrix. With these, we instantiate the weighting choices that satisfy Property (ref), yielding the corollary of Theorem (ref).

We conclude this section by instantiating common feature maps and RL algorithms that satisfy (ref).

Asymptotic Normality of AW-GMM Estimation

In this section, we identify a family of weighting schemes, under which the AW-GMM estimator is asymptotically normal. These weights perfectly stabilize the variance of the empirical influence function in the asymptotic regime. Building on this, we prove strong Gaussian approximation results and characterize the uniform convergence rate. These results are practically appealing, as they enable the construction of uniformly valid confidence regions for parameters of interest over a large class of RL algorithms.

Recall that the AW-GMM estimator $\hat{\theta}_n$ minimizes the empirical loss $\mathcal{L}_n(\theta)$ while the true parameter $\theta^*$ minimizes the expected loss $\mathcal{L}(\theta)$:

The key idea is to relate the estimation error $(\hat{\theta}_n - \theta^*)$ to a martingale difference sequence (MDS), which allows us to apply martingale theory to characterize the asymptotic distribution of $\hat{\theta}_n$.

As established in Section (ref), the estimator $\hat{\theta}_n$ is consistent for $\theta^*$. Since $\theta^*$ lies in the interior of $\Theta$, $\hat{\theta}_n$ will also lie in the interior with high probability. In this case, $\hat{\theta}_n$ satisfies the first-order condition of the empirical loss $\mathcal{L}_n(\theta)$:

Define $\xi_i := \beta_i+J_i\theta^* $. Then after some algebra, we have:

Since $\theta^*$ satisfies the conditional moment equations $\bar{\beta}_i+\bar J_i\theta^*=\mathbf{0}$, the sequence $\{H_i \xi_i\}$ forms an MDS under suitable conditions on the weights $H_i$, as formally established in Section (ref). By linking $\{H_i \xi_i\}$ to the estimation error $(\hat{\theta}_n - \theta^*)$, we derive a strong Gaussian approximation for $\hat{\theta}_n$ and develop corresponding inference results in Section (ref). We then apply these results to common RL algorithms in Section (ref).

Martingale from Adaptive RL Data

We now formalize the martingale difference sequence structure, which plays a central role in our analysis. The true parameter $\theta^*$ satisfies the conditional moment equations, so that $ \bar \beta_i + \bar J_i\theta^* =\mathbf{0}$. Define

where $R_{i,j}$, defined in Lemma (ref), captures the residual component of $Y_i$ after removing the treatment effects from $T_{i,j:L}$ and adding the effects under the evaluation policy $\pi^*_{j:L}$. The sequence $\{\xi_i\}$ forms a martingale difference sequence, as established in the lemma below.

Equation (ref) shows that the estimation error $(\hat{\theta}_n-\theta^*)$ is closely linked to the sum $\sum_{i=1}^nH_i\xi_{i}$. To analyze the asymptotic properties of $\hat{\theta}_n$, it thus suffices to understand the MDS, $\{H_i\xi_i\}$. The following result highlights the key components that characterize its asymptotic behavior.

Proposition (ref) illuminates how weights $H_i$ can be designed to improve the asymptotic properties of the MDS sum $\sum_{i=1}^n H_i\xi_i$, and thus our AW-GMM estimation. A key quantity is the sum of conditional variances, $\sum_{i=1}^n \mathrm{Var}_{i}(H_i\xi_{i})$: when it converges to the identity matrix, this MDS sum is asymptotically normal.

Variance and Covariance under Homoscedasticity.

We now take a closer look at the conditional variance $\mathrm{Var}_{i}(H_i\xi_{i})$, which forms the basis for our subsequent weighting strategy. In full generality, one could design the weights $H_i$ to approximate $\mathrm{Var}_i(\xi_i)^{-1/2}$, so that $\mathrm{Var}_i(H_i \xi_i)$ is close to the identity matrix. However, this approach would require estimating $\mathrm{Var}_i(\xi_i)$, which in turn depends on the state transition dynamics and is thus an impractical task, particularly in high-dimensional continuous state spaces.

To address this challenge, we introduce a homoscedasticity assumption on the outcome process. Under this assumption, the conditional variance matrix $\mathrm{Var}_i(\xi_i)$ becomes block-diagonal: the covariances across different stages (corresponding to the off-diagonal blocks) are zero. As a result, it suffices to design weights conditioning on the realized state $S_{i,j}$ to stabilize the stage-specific conditional variances of $\xi_{i,j}$, thereby eliminating the need to estimate the state transition dynamics.

As shown in Appendix (ref), Assumption (ref) is equivalent to $\mathbb{E}_{i,j}^+[R_{i,j}^2 (\Psi_{i,j} - \bar{\Psi}_{i,j})]=0$, while Lemma (ref) states that $\mathbb{E}_{i,j}^+[R_{i,j} (\Psi_{i,j}-\bar{\Psi}_{i,j})]=0$, reflecting that residuals are uncorrelated with the current treatment conditional on past information. Assumption (ref) strengthens this by requiring the same uncorrelated property for the residual variance. It is satisfied under the additive, rank-preserving version of the structural nested mean model miguel2023causal, but is more general. It implies that, conditional on the past, variance of $R_{i,j}$ is independent of future treatment assignments and is fully explained by exogenous factors. Assumption (ref) holds, for instance, if $R_{i,j}\, {\perp\!\!\!\perp}\, T_{i,j}\mid \mathcal{F}^+_{i,j}$, which is natural if $R_{i,j}$ equals to the potential outcome $\ensuremath{Y}(T_{i,1:j-1}, \pi^*_{j:L})$ in distribution given the past.\footnote{This equality is stronger than Lemma (ref), which requires only equality in expectation.} Given Assumption (ref), we layout the block diagonal structure of the conditional variance $\mathrm{Var}_i(H_i\xi_i)$.

This lemma inspires us to decompose the weight $H_{i,j}$ into two components: one to stabilize $\mathrm{Var}_{i,j}^+(\Psi_{i,j})$ that is known exactly, as it depends on the known behavior policy, and another to stabilize $\mathbb{E}^+_{i,j}\left[ R_{i,j}^2 \right]$ that can be estimated from data.

Strong Gaussian Approximation

We now describe the construction of weights that ensure asymptotic normality of the estimator. This builds on the connection between the MDS $\{\xi_i\}$ and the estimation error $(\hat{\theta}_n - \theta^*)$, as established in Equation (ref):

To establish the asymptotic normality of $\hat{\theta}_n$, it suffices to construct weights such that $\sum_{i=1}^n H_i \xi_i$ is asymptotically normal. Proposition (ref) indicates that this requires stabilizing the conditional variance $\mathrm{Var}_i(H_i \xi_i)$. Lemma (ref) further shows that $\mathrm{Var}_i(H_i \xi_i)$ is block-diagonal, so it is enough to stabilize each stage-specific variance $\mathrm{Var}_i(H_{i,j} \xi_{i,j})$ such that $\sum_{i=1}^n \mathrm{Var}_i(H_{i,j} \xi_{i,j}) \to I$ for each stage $j$. Lemma (ref) shows that:

We thus can split $H_{i,j}$ into two parts: \[ H_{i,j} = \hat{F}_{i,j}^{-\frac{1}{2}}\cdot W_{i,j}, \] where the first part $\hat{F}_{i,j}^{-\frac{1}{2}}$ standardizes $F_{i,j}$ and is estimated using data available up to $\mathcal{F}_{i,j}^+$. Section (ref) provides further discussion on estimating $\hat{F}_{i,j}$; the second part $W_{i,j}$ is designed to satisfy a stabilizing property below that controls the time-varying variance of $\mathrm{Var}^+_{i,j}(\Psi_{i,j})$.

Property (ref) looks similar to Property (ref) but is stronger. In particular, to achieve asymptotic normality, Property (ref)(c) further requires that after weight application, the resulted MDS variance is asymptotically standardized and stabilized around its expectation.

Post-RL Inference.

Note that both $B_n$ and $\widehat{\Xi}_n$ are measurable from the data. Theorem (ref) implies the following Gaussian approximation,

This approximation enables the construction of confidence regions, formalized by the two corollaries below.

Estimating $F_{i,j}$.

We now provide a generic framework for estimating $F_{i,j}$ in a way that satisfies the condition in Equation (ref). Recall that $F_{i,j} := \mathbb{E}^+_{i,j}[R_{i,j}^2]$ is, by definition, a function of $(S_{i,1:j}, T_{i,1:j-1})$. If we further assume that this function is invariant to the behavior policy, then it can be consistently estimated using the realized data. This condition holds, for instance, when the residual $R_{i,j}$ corresponds to the counterfactual outcome $Y(T_{i,1:j-1}, \pi^*_{j:L-1})$, which arises in our application to the Markovian model discussed in Section (ref). In such cases, we denote this function by $f_j$:

Let $\hat{f}_{i,j}$ be a predictor of $f_j$, fitted on the previous $i-1$ episodes and satisfying the estimation rate $i^{-\gamma_F}$ for $\gamma_F\in(0,1)$:

where we introduce the norm $\|\cdot\|_{1,\infty}$ for any $f$ defined on domain $\mathcal{S}^{j_1}\times\mathcal{T}^{j_2}\rightarrow \mathbb{R}$: $\|f\|_{1,\infty} := \mathbb{E}_{s_1\sim P_S}\big[\sup_{(s_{2:j_1},\tau_{1:j_2})\in\mathcal{S}^{j_1-1}\times \mathcal{T}^{j_2}} |f(s_{1:j_1}, \tau_{1:j_2})|\big].$

Then we have Condition (ref) satisfied:

where the last inequality is due to that $x^{-\gamma_F}$ is a concave function with $\gamma_F\in(0,1)$.

Estimating $f_j$ to achieve the rate in (ref) can be addressed by online learning algorithms rakhlin2014online,daskalakis2022fast, and in particular daskalakis2022fast propose a learning algorithm to achieve fast rates of convergence in nonparametric online regression. Later in Section (ref) we show how to estimate $f_{j}$ for high-dimensional Markovian models via Lasso (simpler than proposed algorithm in daskalakis2022fast) and establish the guaranteed convergence rate in (ref).

Normality Weights for Decaying Exploration under Bilinear Features

We now instantiate weight construction on asymptotic normality for common RL algorithms. Similar to Section (ref), we adopt features as instruments and consider the bilinear feature map in Assumption (ref):

where recall that $X_{i,j}$ serves as a shorthand for $\chi_j(S_{i,1:j}, T_{i,1:j-1})$. Under the overlap condition specified in Assumption (ref), we provide weighting choices in the following corollary to achieve strong Gaussian approximation in Theorem (ref).

The choice of $W_{i,j}$ in (ref) requires estimating $V_{i,j}=\mathbb{E}_{i,j}[X_{i,j}X_{i,j}^\top]$. When $j=1$, we have $V_{i,1}=\mathbb{E}\left[X_{i,1}X_{i,1}^\top\right]$, with $\{X_{i,1}\}$ being i.i.d.; thus we can estimate $V_{i,1}$ consistently via the sample-mean estimator and achieve $O(n^{-1/2})$ estimation rate. We hence focus on estimating $V_{i,j}$ for stage $j\geq 2$, where $V_{i,j}$ involves the expectation over $S_{i,j}$ conditional on $(S_{i,1:j-1}, T_{i,1:j-1})$. The randomness in $S_{i,j}$ arises from the state transition process, which is independent of the behavior policy. As a result, we can express $V_{i,j}$ as the output of a function $\nu_j(\cdot)$, with input $(S_{i,1:j-1}, T_{i,1:j-1})$. This function $\nu_j$ depends only on the stage index $j$ but not on the episode index $i$, i.e.,

Estimating $\nu_j$ can be similarly addressed by the online learning literature as we have discussed for estimating $f_j$ in Section (ref), which shall be instantiated with Lasso estimator for the high-dimensional Markovian models in the next section. Below we provide guarantees for generic estimators.

Let $\hat{\nu}_{i,j}$ be a generic predictor of $\nu_j$, fitted on the previous $i-1$ episodes and satisfying the estimation rate $i^{-\gamma_V}$, with $\gamma_V\in(0,1)$:

Above, we reload the norm notation $\|\cdot\|_{1,\infty}$ for any $v$ defined on domain $\mathcal{S}^{j_1}\times\mathcal{T}^{j_2}\rightarrow \mathbb{R}^{d_1\times d_2}$: $\|v\|_{1,\infty} := \mathbb{E}_{s_1\sim P_S}\big[\sup_{(s_{2:j_1},\tau_{1:j_2})\in\mathcal{S}^{j_1-1}\times \mathcal{T}^{j_2}} \|v(s_{1:j_1}, \tau_{1:j_2})\|_{2}\big]$, where $\|\cdot\|_{2}$ denotes the matrix operator norm induced by vector $2$-norm. Then we have Condition (ref) satisfied:

To this end, we provide a full theory with the generic nuisance estimators.

This result is a direct application of Corollary (ref), Equations (ref) and (ref); we omit its proof for brevity.

Application to High-dimensional Markovian Models

In this section, we instantiate the generic inference framework we present in Section (ref) to high-dimensional Markovian models. In particular, we consider the baseline policy of “no treatment” as the evaluation policy, i.e. $0_{1:L}$, where the treatment $0_j$ represents a pre-existing status quo treatment that would have been applied at stage $j$ in the absence of the RL experimentation process. Estimating the value of a no-treatment policy is highly valuable in practical RL scenarios. It enables verification, with statistical confidence, that the deployed RL policy resulted in a statistically significant higher reward compared to consistently applying the status quo treatment at each stage. This allows to test an alternative policy, which introduces a new innovation, product feature, or treatment.

We consider high-dimensional Markovian models that satisfy the high-level assumptions on the data-generating process (DGP) introduced in previous sections. Our focus is on a specific instantiation of the estimation framework in which the features $\Phi_{i,j}$ are used as instruments $\Psi_{i,j}$, and the norm matrix $A$ in the GMM estimation is set to the identity. While our analysis centers on this setting, extensions to more general cases are straightforward. We introduce concrete statistical estimation algorithms for weight construction and provide the corresponding strong Gaussian approximation rates. Finally, we conduct simulations to empirically validate the inferential ability and robustness of our method in this setting.

Data Generating Process

We start by formalizing the DGP for high-dimensional Markovian models. Let $\mathcal{S}, \mathcal{T}$ be the bounded state and treatment spaces. The RL agent sequentially rolls out one episode per unit. Each episode $i$ consists of a length-$L$ trajectory of high-dimensional states, where only a sparse subset of state components influences the final outcome. Specifically, the blip functions depend only on low-dimensional sub-vectors of the state. To facilitate the weight construction discussed in Sections (ref) and (ref), we assume the blip function at each stage $j$ takes a bilinear form and is written as $T_{i,j}\otimes\chi(S_{i,j,\Omega})$ for Lipschitz and bounded $\chi(\cdot)$, where $\Omega$ denotes the set of relevant state coordinates. We use $d_\Omega$ and $d_s$ to denote the dimensions of $S_{i,j,\Omega}$ and $S_{i,j}$, respectively.

For each episode $i$, the behavior policy $\pi^{\text{obs}}_i$ is decided based on the previous $i-1$ episodes. This policy $\pi^{\text{obs}}_i$ is assumed to be known and assigns the treatment via $T_{i,j} := \pi^{\text{obs}}_i(S_{i,j}, \zeta_{i,j})$, where $\zeta_{i,j}$ is an i.i.d. bounded noise term. We consider the behavior policy satisfies Assumption (ref)(a) such that:

for an exploration rate $\alpha$. State transitions are governed by unknown matrices $(A_j, B_j, M_j)$, which capture the effects of the current treatment, the current state, and the initial state, respectively, on the next state. Let $\eta_{i,j}$ denote the i.i.d. mean-zero bounded noise associated with the state transitions. The final outcome $Y_i$ is modeled as a sparse linear function of the last treatment, the last state, and the initial state, perturbed by an i.i.d. bounded noise term $\epsilon_i$. Model (ref) summarizes the DGP.

Given the DGP outlined in Model (ref) , we can expand the outcome $\ensuremath{Y}$ recursively:

Our goal is to estimate the structural parameters $\{\theta_0^*, \theta_1^*,\dots, \theta_L^*\}$, where $\theta_0^*:=\mathbb{E}[Y(0_{1:L})]=\beta_1^\top\mathbb{E}[S_{i,1}] + \kappa_1^\top\mathbb{E}[ S_{i,1}]$ represents the expected outcome under the baseline policy, and $\theta_{1:L}^*$ reflect the stage-wise dynamic treatment effects in Model (ref) outlined in (ref).

Constructing the Weights

We now provide details on the weight construction for Model (ref). With the bilinear blip functions, Corollary (ref) provides analytical weight construction to achieve consistency. To achieve the strong Gaussian approximation, Corollary (ref) shows that we need to estimate functions $f_j$ and $\nu_j$ to satisfy Conditions (ref) & (ref) respectively. Recall that in general settings, $f_{j}$ and $\nu_{j}$ can be estimated by online learning algorithms, while here under the high-dimensional Markovian models, this estimation procedure can be greatly simplified.

Estimating $f_{j}$.

We first show that Model (ref) satisfies the below homoscedasticity property, which satisfies Assumptions (ref) & (ref) and makes estimating $f_{j}$ a regular estimation problem.

With this property, we can decompose $F_{i,j}$ into two parts:

Note that $G_{i,j}$, when conditioning on $(S_{i, 1:j}, T_{i,1:j-1})$, does not depend on the behavior policy. Thus we can view $G_{i,j}$ as a function $g_j$ of $(S_{i,1:j}, T_{i,1:j-1})$, which only depends on the stage index $j$ but not the specific episode index $i$, that is,

where the last equation is due to that $R_{i,j}=\beta_j^\top S_{i,j} + \kappa_j^\top S_{i,1} + \epsilon_{i,j}$ by Equation (ref). We can estimate $g_j$ by regressing $\{R_{i,j}\}$ on $\{(S_{i,1}, S_{i,j})\}$ to get $\hat{g}_{i,j}$; we then estimate $\sigma_{j}^2$ by its sample variance $\hat\sigma_{i,j}^2$ up till episode $i$; finally we set $\hat{f}_{i,j}(\cdot):=\mathrm{Clip}_{[\sigma^2, M^2]}\left(\hat{g}_{i,j}(\cdot)^2 + \hat\sigma_{i,j}^2\right)$, where $(\sigma^2,M^2)$ denotes the specified range for $F_{i,j}$. Here, for any real numbers $a < b$ and any $x \in \mathbb{R}$, the clipping operator is defined as:

Algorithm (ref) summarizes the estimation process. Note that we cannot observe the true residual $R_{i,j} =Y_{i}-\sum_{j'>j}\Phi_{i,j'}^\top\theta^*_j $, which requires the knowledge of the true structure parameter $\theta^*$. However, we can estimate it using an approximated $\hat{\theta}$ (for example, the consistent estimate from Section (ref)). With that, we obtain the estimation rate of $\hat{f}_{i,j}$ as $\|\hat f_{i,j}-f_j\|_{1,\infty}=O(i^{-\frac{1-L\alpha}{2}})$; see details in Appendix (ref).

\setcounter{algocf}{0}

Estimating $\nu_{j}$.

We now estimate \[\nu_j(s_{1:j-1},\tau_{1:j-1}):=\mathbb{E}\left[\chi(S_{i,j,\Omega})\chi(S_{i,j,\Omega})^\top\mid S_{i,1:j-1}=s_{1:j-1}, T_{i,1:j-1}=\tau_{1:j-1}\right].\] Note that when $j=1$, we have $\nu_{1}\equiv\mathbb{E}\left[\chi(S_{i,1,\Omega})\chi(S_{i,1,\Omega})^\top\right]$, with $\{S_{i,1}\}$ being i.i.d.; thus we can estimate $\nu_{1}$ consistently via the sample-mean estimator and achieve $O(n^{-1/2})$ estimation rate. We hence focus on estimation of $\nu_{j}$ for stage $j\geq 2$.

Let $h_{j}(S_{i,1:j-1}, T_{i,1:j-1})$ denote the conditional expectation of $S_{i,j}$ when restricted to coordinates $\Omega$ (since only those coordinates enter into $\psi_j$ and contribute to $V_j$):

Then $\nu_j$ can be induced by $h_j$: $ \nu_j(\cdot) =~ \mathbb{E}_{i,j}\left[\chi\left(h_{j}(\cdot) + \eta_{i,j,\Omega}\right)\chi\left(h_{j}(\cdot) + \eta_{i,j,\Omega}\right)^\top\right], $ for i.i.d. exogenous noise $\eta_{i,j,\Omega}$ during the state transition. We therefore first regress $\{S_{i,j,\Omega}\}$ on $\{S_{i,1}, S_{i,j-1},T_{i,j-1}\}$ to get estimate $\hat{h}_{i,j}$, based on which we obtain estimate $\hat \nu_{i,j}$, as summarized in Algorithm (ref). Note that in Algorithm (ref), we also apply the clipping operator when obtaining $\hat{\nu}_{i,j}$ to ensure its eigenvalues lie within the specified range for the true matrix $V_{i,j}$. Here, For any real numbers $a < b$ and a symmetric matrix $X = U \Lambda U^\top \in \mathbb{R}^{d \times d}$, where $U$ is an orthogonal matrix and $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_d)$ is the diagonal matrix of eigenvalues, the clipping operator is defined as:

where recall $\text{Clip}_{[a,b]}(\lambda) := \max(\min(\lambda, b), a)$ for any scalar $\lambda \in \mathbb{R}$ as defined (ref). The estimation rate of $\hat{\nu}_{i,j}$ satisfies $\|\hat\nu_{i,j}-\nu_j\|_{1,\infty}=O(i^{\frac{-\gamma(1-2\alpha)}{2}})$ (proof deferred to Appendix (ref)).

Putting everything together.

With the estimated $\hat{f}_{i,j}$ and $\hat{\nu}_{i,j}$ in the previous sections, we are ready to construct weights to achieve strong Gaussian approximation in Corollary (ref). We summarize the full steps in Algorithm (ref) and provide the strong Gaussian approximation rate in Corollary (ref).

Numerical Experiments

We finally present empirical evidence supporting our method's effectiveness in scenarios with both correct and mis-specified feature mapping function $\phi_j$. Our findings illustrate that applying weights enhances outcomes in both cases: under correct specification, it empirically confirms our method's consistency and asymptotic normality across all structural parameters. In cases of mis-specification, it notably improves the accuracy of estimating the evaluation policy value.

We study a two-stage high-dimensional Markovian model (with $L=2$), as outlined in Model (ref), focusing on binary treatment scenarios. Data collection is performed by an $\epsilon$-greedy episodic RL agent, which sequentially rolls out an episode for each unit. The agent's behavior policy undergoes batch updates, with each batch including $100$ units; the exploration amount $\epsilon$ decays at a polynomial rate such that for any given batch index $b$, the exploration $\epsilon$ is set to be $b^{-0.5}/2$, satisfying Assumption (ref). Throughout the experiment, a total of $5000$ units are collected.

We consider high dimensional states and low dimensional features, with sparse linear models for both the state transition and the final outcome. Specifically, the state space is in $\mathbb{R}^{10}$ with $S_{i,j}= (S_{i,j,1},\dots, S_{i,j,10})$; only the first coordinate is informative, with others being noises. We focus on GMM-estimations using four different weighting schemes:

Note that the Oracle weights are infeasible, as they require knowledge of the true data-generating process. Conversely, the Naive, Consistent, and Feasible weights can either be directly computed or estimated from the data. Below we show that Feasible weighting scheme performs comparably to the Oracle and significantly outperforms the other two in achieving asymptotic normality for post-RL inference; the \textsc{Feasible} also demonstrates robustness under misspecification.

Estimation and Inference Validity.

We first consider cases with correctly specified feature mapping. Define the true feature mapping as $\phi_j(S_{i,1:j}, T_{i,1:j})=T_{i,j}\cdot (S_{i,j,1}, S_{i,j,1}^2)\in\mathbb{R}^2$. Our estimand, $\theta^*=(\theta_0^*, \theta_{1,1}^*, \theta_{1,2}^*, \theta_{2,1}^*, \theta_{2,2}^*)$, includes the evaluation policy value $\theta_0^*$ and structure parameters $(\theta_{1,1}^*, \theta_{1,2}^*)$ and $(\theta_{2,1}^*, \theta_{2,2}^*)$, indicating the effect of the treatment at each stage. Figure (ref) shows that estimations with Oracle and Feasible weights achieve nominal coverage, contrasting with the Naive no-weighting or Consistent weights, which either have low coverage or are overly conservative. In particular, the Naive estimator's confidence intervals for the policy value $\theta_0^*$ and the first-stage treatment effect $(\theta_{1,1}^*,\theta_{1,2}^*)$ even widen with increased sample size. Figure (ref) further shows that the studentized statistics of the AW-GMM estimators, derived from Eq.(ref), under Oracle and \textsc{Feasible} weights conform to asymptotic normality, unlike those from other weighting schemes. Furthermore, Figure (ref) shows that while all estimators yield similar results for the stage-wise treatment effect estimation, those with \textsc{Oracle} and \textsc{Feasible} weights achieve higher accuracy in estimating the policy value $\theta_0^*$ with smaller MSE and bias.

Robustness under Misspecification.

Transitioning to cases of mis-specification, we consider a true feature mapping defined as $\phi_j(S_{i,1:j}, T_{i,1:j})=\exp(T_{i,j}S_{i,j,1}/2) -1$, unknown to the AW-GMM estimators. These estimators then use polynomial approximations with degree $d\in[1:5]$ instead of the true $\phi_j$, i.e., we use $\hat{\phi}^{(d)}_j(S_{i,1:j}, T_{i,1:j})=T_{i,j}(S_{i,j,1}, \dots, S_{i,j,1}^d)\in\mathbb{R}^d$ as the approximated feature mapping for each degree $d$. Under this scenario, the stage-wise structural parameters lose their causal interpretation, yet estimating the evaluation policy value $\theta_0^*$ remains relevant. Figure (ref) shows that AW-GMM estimations under Feasible weights outperform those with Naive no-weighting or Consistent weights, demonstrating better coverage, narrower confidence intervals, and reduced MSE and bias. Interestingly, with Feasible weights, improvements in inference validity (coverage) and efficiency (confidence interval length) plateau for approximation degrees $d \geq 3$, while estimation quality (MSE and bias) slightly deteriorates at $d=5$. The decrease in estimation accuracy at higher approximation degrees is due to the requirement to estimate more structural parameters from the same sample size, which instead complicates the estimation process. This observation highlights a balance between inference and estimation: without knowing the precise feature mapping, selecting an approximation with appropriate complexity is crucial in optimizing both estimation accuracy and inferential robustness.

Proofs of Main Lemmas

Proof of Lemma (ref)

We follow the proof pattern for Lemma 6 in lewis2020double. Note that for the observed $\ensuremath{Y}$, we always have $\ensuremath{Y}\equiv \ensuremath{Y}(T_{1:L})$. For any stage $j$, we have

where (i) is by Assumption (ref), (ii) is by the definition of blip function. Moreover, note that:

where (iii) follows since the counterfactual outcomes are independent of the value of $T_j$ conditional on $S_{1:j}, T_{1:j-1}$ by Sequential Conditional Exogeneity Assumption (ref). Subtracting the two equalities, and by the definition of $\Phi_j$, we derive that:

Then we have

where (i) uses a telescoping sum. Rearranging the above, we have

Proof of Lemma (ref)

We adapt the proof pattern of Lemma 7 in lewis2020double to the RL data.

Moreover, note that by Lemma (ref) and the definition of $R_{i,j}$, we have:

Moreover, by Assumption (ref), we have:

Combining the last three equations, we conclude that:

where the last equality holds, since by definition $\bar{\Psi}_{i,j}:=\mathbb{E}_{i,j}^+[\Psi_{i,j}]$.

Proof of Lemma (ref)

We show that $\sum_{i=1}^n H_i\xi_{i}$ is a sum of martingale difference sequence adapted to filtration $\{\mathcal{F}_i\}$. By Lemma (ref) and the definition of $\xi_{i,j}$, we have for any $j\in \{0,\dots,L\}$

where $(i)$ follows by construction that $H_{i,j}$ is measurable with respect to $\mathcal{F}^+_{i,j}$ and $(ii)$ follows by Lemma (ref). Therefore we have:

implying $ \mathbb{E}_i\left[H_{i}\xi_{i} \right]=0$.

Proof of Lemma (ref)

We first prove that $\mathbb{E}[R_{i,j}^2\mid T_{i,j}, \mathcal{F}_{i,j}^+]$ does not depend on $T_{i,j}$ such that $\mathbb{E}[R_{i,j}^2\mid T_{i,j}, \mathcal{F}_{i,j}^+]=\mathbb{E}[R_{i,j}^2\mid \mathcal{F}_{i,j}^+]=\mathbb{E}_{i,j}^+[R_{i,j}^2]$. Define $\Bar{R}_{i,j}:=\mathbb{E}[R_{i,j}\mid T_{i,j}, \mathcal{F}^+_{i,j}]$. We have:

By Assumption (ref), $\mathrm{Var}(R_{i,j}\mid T_{i,j}, \mathcal{F}_{i,j}^+) \equiv\mathrm{Var}(R_{i,j}\mid \mathcal{F}_{i,j}^+)$ does not depend on $T_{i,j}$. Moreover, Appendix (ref) proves that $\Bar{R}_{i,j}$ is independent of $T_{i,j}$ and thus $\Bar{R}_{i,j}=\mathbb{E}_{i,j}^+[R_{i,j}]$. We therefore have

which does not depend on $T_{i,j}$.

\paragraph{Part (a).} Consider the conditional covariance. Without loss of generality, assume $j_1<j_2$. Since $H_{i,j}$ are measurable with respect to $\mathcal{F}^+_{i,j}$, we can write:

Note that:

Thus:

Therefore, we have:

\paragraph{Part (b).} By the definition of $\xi_{i,j}$ and the fact that the weights $H_{i,j}$ are measurable in $\mathcal{F}_{i,j}^+$:

Therefore,

Proof of Theorem (ref)

We hereby show that GMM estimator $\hat{\theta}_n\in\argmin_{\theta\in\Theta}\mathcal{L}_n(\theta)$ converges to the true parameter $\theta^*\in\argmin_{\theta\in\Theta}\mathcal{L}(\theta)$, where

In Appendix (ref), we show that uniformly across $\theta\in\Theta$, we have:

Then for $\hat\theta_n$ that minimizes the loss $\mathcal{L}_n(\theta) = \|m_n(\theta)\|_A^2$, we have:

Therefore,

On the other hand, we have:

With (ref), Appendix (ref) applies the induction method to invert $\frac{1}{n}\sum_{i=1}^n H_i \bar{J}_i$ and show that for each $j\in[0,L]$, it holds that:

Concluding the proof.

Uniform Convergence

We now show:

We have:

where for a symmatric matrix $M$, we use $\lambda_{\max}(M)$ to denote its largest eigenvalue and use $\|M\|_2$ to denote its operator norm induced by vector 2-norm. Appendix (ref) shows that $\mathbb{E}[C^2]=O(Ln^{\alpha_2-1})$, and Appendix (ref) shows that $\mathbb{E}[D^2]=O(L^2n^{\alpha_2-1})$. Therefore, \[ \mathbb{E}[I^2]\lesssim \mathbb{E}[C^2]+ \mathbb{E}[D^2] = O(L^2n^{\alpha_2-1}), \] proving (ref).

Bounding term C.

We show:

Note that $\{H_i(\beta_i-\bar \beta_i)\}$ is a martingale difference sequence adapted to $\{\mathcal{F}_i\}$. This is because each $j$-th component \[ \mathbb{E}_i[H_{i,j}(\beta_{i,j}-\bar \beta_{i,j})] = \mathbb{E}_i \mathbb{E}_{i,j}^+[H_{i,j}(\beta_{i,j}-\bar \beta_{i,j})] = \mathbb{E}_i \left[H_{i,j}\mathbb{E}_{i,j}^+[(\beta_{i,j}-\bar \beta_{i,j})]\right]=0, \] by the definition $\bar \beta_{i,j}=\mathbb{E}_{i,j}^+[\beta_{i,j}]$ and $H_{i,j}$ adapted to $\mathcal{F}_{i,j}^+$. We have

where (i) is due to that $\{H_i(\beta_i-\bar \beta_i)\}$ is a martingale difference sequence. Now consider each diagonal entry $j$. We have:

where the last equality is by Property (ref)(c). We thus have (ref) hold.

Bounding term D.

We next show:

Similarly, $\{H_i(J_i-\bar J_i)\}$ is a martingale difference sequence adapted to $\{\mathcal{F}_i\}$. For any given $\theta\in\Theta$ with $\Theta$ being bounded, we have

where $\Lambda_{i,j,k}:=\mathrm{Tr}\left(\mathbb{E}_{i,j}^+\left[H_{i,j}\, (J_{i,j,k} -\bar{J}_{i,j,k})(J_{i,j,k} -\bar{J}_{i,j,k})^\top H_{i,j}^\top \right]\right)$; (i) uses the boundedness of $\theta$ and that a matrix $\ell_2$ norm is bounded by its Frobenius norm; (ii) uses that $\{H_i\, (J_i -\bar{J}_i)\}_{i=1}^n$ is a martingale difference sequence; (iii) is due to that $H_i$ is a block-diagonal matrix and $J_i$ is a block upper-triangular matrix. Now let's look into (ref) and analyze $\Lambda_{i,j,k}$.

where (i) is due to Lemma (ref), boundedness of $\|\Phi_{i,k}\|_2$, and that if positive semi-definite matrices $A\preceq B$, then $\mathrm{Tr}(CAC^\top)\leq \mathrm{Tr}(CBC^\top)$ for any matrix $C$.

Continuing (ref), we have:

where (i) is by Property (ref)(c). Thus we have (ref) hold.

Induction Step

We now bound the estimation error $(\hat\theta_n - \theta^*)$. Equation (ref) shows that the normalized error satisfies $ \mathbb{E}[\|\bar B_n (\hat{\theta}_n-\theta^*) \|_{2}^2]=O(n^{\alpha_2-1})$, for $ \Bar B_{n} := -n^{-1}\sum_{i=1}^nH_i \bar{J}_i$. Note that $\Bar{B}_n$ is an upper block-triangular matrix. We will show that, for each diagonal block $\Bar{B}_{n,j,j}$, the matrix $\Bar{B}_{n,j,j}^\top \Bar{B}_{n,j,j}$ is invertible with high probability. This allows us to invert $\Bar{B}_{n,j,j}^\top \Bar{B}_{n,j,j}$ to bound the estimation error $(\hat{\theta}_{n,j} - \theta_j^*)$ in a backwards manner.

In particular, we apply induction method to prove the result. We start by introducing a few notations. With $\Bar B_{n} = -n^{-1}\sum_{i=1}^nH_i \bar{J}_i$, define $\Bar B_{n,j,j'}=n^{-1}\sum_{i=1}^n H_{i,j}\mathbb{E}_{i,j}^+\left[(\Psi_{i,j}-\bar{\Psi}_{i,j})\Phi_{i,j'}^\top\right]$ as its $(j,j')$ block. Similarly, define $ B^0_{n,j,j'}=n^{-1}\sum_{i=1}^n \mathbb{E}_{i,j}\left[H_{i,j}(\Psi_{i,j}-\bar{\Psi}_{i,j})\Phi_{i,j'}^\top\right]$. Define their difference $\delta_{n,j,j'}:=B^0_{n,j,j'}-\Bar B_{n,j,j'}$. Appendix (ref) shows that:

By Property (ref)(a), we have $\forall l\in [0:L]$, $(B^0_{n,j,j})^\top B^0_{n,j,j} \succeq c_1^2 n^{-\alpha_1}I$, where $c_1$ is a constant introduced in Property (ref). The regularity of $\bar B_{n,j,j}$ inherits from that of $B^0_{n,j,j}$. Define the event $\mathcal{E}_n$ as:

Appendix (ref) shows that $\mathcal{E}_n $ happens with high probability:

We now apply the induction step to bound the estimation error $\hat\theta_n - \theta^*$.

\paragraph{Base case: $j=L$.} We have:

where (i) is because $B_n$ is a block upper triangular matrix and (ii) is by (ref) and (ref). This result also yield the $L^1$-norm convergence:

\paragraph{Induction step.} Now consider $j<L$ recursively. Assume the induction hypothesis holds that for $j'=j+1, \dots, L$, such that:

Let's first prove the $L^2$-norm convergence. We have

where (i) is by triangular inequality and Cauchy-Swarchz inequality, (ii) is because $\Bar B_n$ is a block upper triangular matrix and (iii) is by (ref), (ref), (ref), and (ref).

We then look into the $L^1$-norm convergence by having:

where (i) uses that $\bar B_n$ is a block upper triangular matrix, (ii) uses (ref), triangular inequality, and (ref), (iii) uses Cauchy-Schwartz inequality and that a matrix $\ell_2$ norm is bounded by its Frobenius norm, (iv) uses the induction assumption (ref), and (v) uses that $\alpha_2\leq \alpha_1\leq 1/L$ as specified in Property (ref).

Asymptotic regularity of $\Bar{B}_n$.

Note that $\Bar{B}_n$ is block upper triangular, where the column sizes of blocks correspond to the decomposition of $\theta=(\theta_0, \theta_1, \dots, \theta_L)$. For any $j\leq j'$ we have:

Define the matrix $B_n^0$ as follows: for each of its $(j,j')$-th block with $j\leq j'$, define

First, notice that the singular values of the $j$-th diagonal block $B^0_{n,j,j}$ are lower-bounded when Property (ref)(a) holds. For the $0$-th diagonal block corresponding to $\theta_0^*$, we have $B^0_{n,0,0}= n^{-1}\sum_{i=1}^n \mathbb{E}_{i,0}[H_{i,0}] =1\geq c_1 n^{-\frac{\alpha_1}{2}}$ for $H_{i,0}=1$ and large $n$.

Next we shall show that for $j\leq j'$, we have

and that the event $\mathcal{E}_n$ happens with high probability:

\paragraph{Asymptotic neglibility of $\delta_{n,j,j'}$.} For a matrix $A$, we use $\mathrm{Tr}(A)$ to denote its trace, i.e. the sum of its diagonal entries. With $j\leq j'$, by the definition of the Frobenius norm, we have

where (i) uses the fact that $\delta_{n,j,j'}$ is a sum of martingale difference sequence, (ii) uses Lemmas (ref) & (ref), (iii) uses that $\|\Phi_{i,j'}\|_2$ is bounded, Lemma (ref) and Lemma (ref), and (iv) uses Property (ref)(c).

\paragraph{Uniform bound of $B^0_{n,j,j'}$.} With $j\leq j'$, we have

where (i) is by triangular inequality, (ii) is by Cauchy-Schwartz inequality, (iii) is by the assumption that $\|\Phi_{i,j'}\|_2$ is bounded, Lemma (ref) and Lemma (ref).

\paragraph{Uniform bound of $\Bar B_{n,j,j'}$.} With $j\leq j'$, we have

where (i) is by triangular inequality and Cauchy-Schwartz inequality, (ii) is due to the fact that $\Phi_{i,j'}, \Psi_{i,j},\bar{\Psi}_{i,j}$ are bounded such that \[ H_{i,j}\mathbb{E}_{i,j}^+[(\Psi_{i,j}-\bar{\Psi}_{i,j})\Phi_{i,j'}^\top] \mathbb{E}_{i,j}^+[\Phi_{i,j'}(\Psi_{i,j}-\bar{\Psi}_{i,j})^\top] H_{i,j}^\top\precsim H_{i,j}H_{i,j}^\top; \] then with Lemma (ref) we have (ii).

\paragraph{High probability event $\mathcal{E}_n$.} We finally show that $\mathbb{P}(\mathcal{E}^c_n)=O(Ln^{\alpha_1+\alpha_2-1})$. When $\mathcal{E}_n$ does not happen, there exists a $j\in\{0,1,\dots, L\}$ and an eigenvector $x$ of $\Bar B_{n,j,j}^\top \Bar B_{n,j,j}$, with $\|x\|_2=1$ such that $\Bar B_{n,j,j}^\top \Bar B_{n,j,j} x = \lambda x$ and $\lambda<c_1^2n^{-\alpha_1}/4$. In that case:

Since $\delta_{n,j,j}=\Bar B_{n,j,j}-B_{n,j,j}^0$, we have that:

where (i) is by Lemma (ref), and (ii) is because the minimum eigenvalue value of $(B_{n,j,j}^0)^\top B_{n,j,j}^0$ is at least $c_1^2n^{-\alpha}$ by Property (ref)(a) and the vector $x$ is unit-norm. Rearranging yields:

Thus we have, by Markov's inequality:

Proof of Theorem (ref)

For notation convenience, we write $F_{i,j}:=f_i(S_{i,1:j}, T_{i,1:j-1})$ and $\hat{F}_{i,j}:=\hat{f}_{i,j}(S_{i,1:j}, T_{i,1:j-1})$. Also, we let $F_i:=(F_{i,1},\dots, F_{i,L})$, and $\hat{F}_i=(\hat{F}_{i,1}, \dots, \hat{F}_{i,L})$. Define the variance estimate: $\widehat{\Xi}_n=n^{-1}\sum_{i=1}^n \mbox{diag}\left\{W_{i,j}\mathrm{Var}^+_{i,j}(\Psi_{i,j})W_{i,j}^\top\right\}$.

Let $\hat{\theta}_n$ be the AW-GMM estimator defined in (ref). With assumptions in Theorem (ref), the weights $H_{i,j}$ satisfy Property (ref) with constants $(\alpha_1, 0)$. We can then invoke Theorem (ref) and get:

Thus by Markov inequality, we have for any positive $\delta$:

With $\theta^*$ lying in the interior of $\Theta$ and $\Theta$ being bounded, for $n$ larger than a constant, with probability at least $1-O(n^{\frac{(L+1)\alpha_1-1}{2}})$, $\hat{\theta}_n$ lies in the interior of $\Theta$. Define $\mathcal{E}_{int}$ as the event that $\hat{\theta}_n$ lies in the interior of $\Theta$. Then for $n$ larger than a constant, $\mathbb{P}(\mathcal{E}_{int}^c)=O(n^{\frac{(L+1)\alpha_1-1}{2}})$.

Now let's establish the strong Gaussian approximation when event $\mathcal{E}_{int}$ happens. With $\hat{\theta}_n$ being in the interior of $\Theta$, we can invoke the first-order condition of the GMM solution:

Appendix (ref) shows that the sum of MDS $\frac{1}{\sqrt{n}}\sum_{i=1}^n H_i\xi_{i}$ can be approximated uniformly by a Gaussian random variable $Z_\xi\sim \mathcal{N}\left(0, \widehat{\Xi}_n\right)$, such that

where $\mathcal{C}$ is the set of all convex subsets in $\mathbb{R}^{1+dL}$.

Therefore for any convex set $\bar C$,

Now for the convex set $\bar{C}$, define the induced set $\tilde{C}$ as follows:

This definition yields:

Meanwhile, it's straightforward to see that $\tilde{C}$ is also a convex set; thus by the strong Gaussian approximation in (ref), we have

Combining (ref), (ref), (ref), we have for any convex set $\bar{C}\in\mathcal{C}$:

We finally show that the normalizing matrix $B_n^\top AB_n$ is invertible with high probability. Define event $\mathcal{E}^M_n$ for $B_n=-n^{-1}\sum_{i=1}^n H_iJ_i$:

Appendix (ref) shows that $\mathcal{E}^M_n$ happens with probability at least $1-O(Ln^{\alpha_1-1})$. When $\mathcal{E}^M_n$ happens, each $B_{n,j,j}$ has full column rank. Since $B_n$ is a block upper triangular matrix, $B_n$ also has full column rank. As a result, $B_n^\top A B_n$, which has the same rank as $B_n$ (see Lemma (ref)), is thus invertible.

Now for any convex set $C\in\mathcal{C}$, define the induced set

which is also a convex set by construction. We have:

This concludes our proof for Theorem (ref).

Strong Gaussian Approximation of $n^{-\frac{1}{2}}\sum_{i=1}^n H_i\xi_i$

We show the strong Gaussian approximation results of $n^{-\frac{1}{2}}\sum_{i=1}^n H_i \xi_i$. To do it, we follow three steps:

Collectively, we shall be able to connect $n^{-1}\sum_{i=1}^n H_i\xi_i$ with $\mathcal{N}(0, \widehat{\Xi}_n)$, which achieves the inferential result we target.

Step 1: Connect $n^{-\frac{1}{2}}\sum_{i=1}^n H_i\xi_i$ with $\mathcal{N}(0, \Sigma_n)$.

Recall that $\{ H_i\xi_i\}_{i=1}^n$ is a martingale difference sequence. To characterize its asymptotic behavior, we shall leverage the following proposition.

We now write out the terms in Proposition (ref) regarding $\{ H_i\xi_i\}_{i=1}^{n}$:

\paragraph{Regularity of $\Sigma_{n}$.} We first lower bound the eigenvalues of $\Sigma_{n}$. Note that by construction, $\Sigma_{n}$ is a symmetric block-diagonal matrix, and thus we only need to look at its $j$-th diagonal block $\Sigma_{n,j,j}$. Define $\Xi_{n} := n^{-1}\mbox{diag}_j\left\{\mathbb{E}\left[\sum_{i=1}^n W_{i,j}\mathrm{Var}^+_{i,j}(\Psi_{i,j})W_{i,j}^\top\right]\right\}$, and so $\Xi_n$ is also a block-diagonal and symmetric matrix. We have

Thus we have

Together with Property (ref)(c) that says that $\Xi_{n}\succeq c_6\cdot I$ and Lemma (ref), we have that

Therefore, for large enough $n$, $\Sigma_{n} \succeq \frac{c_2}{2}I$ is invertible. Moreover, we have:

\paragraph{Magnitude of $\Omega_{n}$.} Note that $\Omega_{n}$ by definition is block-diagonal and symmetric.

\paragraph{Regularity of $\beta_2$.} We move onto discussing $\beta_2$ defined in Proposition (ref) and have that

where $\kappa_4$ is the fourth-moment of a standard normal random variable, (i) uses that \[ \frac{F_{i,j}}{\hat{F}_{i,j}} W_{i,j}\mathrm{Var}_{i,j}^+(\Psi_{i,j})W_{i,j}^\top\preceq\frac{M^2}{\sigma^2}\mathrm{Tr}(W_{i,j}\mathrm{Var}_{i,j}^+(\Psi_{i,j})W_{i,j}^\top)\cdot I \precsim I, \] and that \[ \mathbb{E}\left[\| H_i\xi_{i}\|^2\right]=\mathrm{Tr}(\mathbb{E}[v_i])=\mathrm{Tr}\left(\mathbb{E}\left[\mbox{diag}_j\left\{\frac{F_{i,j}}{\hat{F}_{i,j}} W_{i,j}\mathrm{Var}_{i,j}^+(\Psi_{i,j})W_{i,j}^\top\right\}\right]\right)=O(Ld'); \] and (ii) uses Jensen's inequality.\footnote{For any sequence $a_1,\ldots, a_K$: $(\sum_{i=1}^K a_i^2)^2 = K^2 \left(\frac{1}{K}\sum_{i=1}^K a_i^2\right)^2 \leq K\sum_{i=1}^K\left( a_i^2\right)^2$.}

Now we consider $\mathbb{E}\left[\| H_i\xi_i\|_2^4\right]$.

where the inequality is by Cauchy-Schwartz inequality. Note that $\hat{F}_{i,j}\geq \sigma^2$, and both $R_{i,j}^4$ and $\|\Psi_{i,j}-\bar{\Psi}_{i,j}\|_2^2$ are bounded by some finite constant, thus we have $(\Psi_{i,j}-\bar{\Psi}_{i,j})(\Psi_{i,j}-\bar{\Psi}_{i,j})^\top \precsim I_d$. Hence,

where (i) uses that if $A,B$ are positive semi-definite, we have $\mathrm{Tr}(AB)\leq\mathrm{Tr}(A)\mathrm{Tr}(B)$, and (ii) uses Property (ref)(c).

Combining (ref) and (ref), we have that

where the last equality is by Property (ref)(b).

\paragraph{Applying Proposition (ref).} Combining (ref), (ref), (ref), We have

Choose $\eta= n^{\frac{5-2\gamma'}{10}}$ where $\gamma'\in(0,\min(\gamma_F, \alpha_3)]$. Continuing the above, we have

Collectively, we have

For the remaining steps $2-3$, we need to invoke the following lemma.

Step 2: Connect $\mathcal{N}(0, \Sigma_n)$ with $\mathcal{N}(0, \Xi_n)$.

Note that by (ref), we have that $\left\|\Sigma_n - \Xi_n\right\|_{Frob}^2=O(Ln^{-2\gamma_F})$. Moreover, under Property (ref)(c), $\Xi_n \succeq c_2 I $ and thus is invertible; also $\Sigma_n\succeq \frac{c_2}{2}I$ for large $n$ by (ref). Consider

Let $\{\lambda_i\}$ be the eigenvalues of $\Xi_n^{-1}\Sigma_n-I $, we have

Applying Lemma (ref), we thus have

Step 3: Connect $\mathcal{N}(0,\Xi_n)$ with $\mathcal{N}(0, \widehat{\Xi}_n)$.

We first show that event $\widehat{\mathcal{E}}_n=\left\{ \widehat{\Xi}_n\succeq \frac{c_6}{2} I\right\}$ happens with high probability. Define $\overline{\Xi}_n$ as a block-diagonal matrix with its $j$-th diagonal block as, \[ \frac{1}{n}\sum_{i=1}^{n} \mathbb{E}_{i}[W_{i,j}\mathrm{Var}^+_{i,j}(\Psi_{i,j})W_{i,j}^\top]. \] Note that $\widehat{\Xi}_n-\overline{\Xi}_n$ is the sum of martingale difference sequence. We have

Thus \[ \mathbb{E}\left[ \left\|\widehat{\Xi}_n-\overline{\Xi}_n\right\|_{2} \right]\leq \mathbb{E}\left[ \left\|\widehat{\Xi}_n-\overline{\Xi}_n\right\|_{Frob} \right] \leq \mathbb{E}\left[ \left\|\widehat{\Xi}_n-\overline{\Xi}_n\right\|_{Frob}^2 \right]^{1/2}=O\left(L^{\frac{1}{2}}n^{-\frac{1}{2}}\right). \] Therefore, \[ \mathbb{E}\left[ \left\|\widehat{\Xi}_n - {\Xi}_n\right\|_{2} \right]\leq \mathbb{E}\left[ \left\|\widehat{\Xi}_n-\overline{\Xi}_n\right\|_{2} \right] + \mathbb{E}\left[ \left\|\overline{\Xi}_n - {\Xi}_n\right\|_{2} \right] = O(L^{\frac{1}{2}}n^{-\min(\frac{1}{2}, \alpha_3)}), \] where we use the result that $ \mathbb{E}[\|\overline{\Xi}_n-\Xi_n\|_2] = O(n^{-\alpha_3})$, which holds as proven below:

On the other hand, Property (ref)(c) says that $\Xi_n\succeq c_2\cdot I$. Applying Lemma (ref), we have that: \[ \mathbb{P}(\widehat{\mathcal{E}}_n^c) = O(L^{\frac{1}{2}}n^{-\min(\alpha_3,\frac{1}{2})}). \]

Under Property (ref)(c), we have that $\Xi_n \succeq c_2 I $ and thus is invertible. Now conditioning on event $\widehat{\mathcal{E}}_n$ (and thus $\widehat{\Xi}_n$ is invertible), we have

We overload notations and let $\{\lambda_i\}$ be the eigenvalues of $\Xi_n^{-1}\widehat{\Xi}_n-I $, we have

Applying Lemma (ref), we have that, conditioning on event $\widehat{\mathcal{E}}_n $ happens,

where both sides are functions of the random matrix $\widehat{\Xi}_n$, and the total variation distance on the left-hand-side is defined conditional on $\widehat{\Xi}_n$, which serves as the covariance matrix of the Gaussian distribution $\mathcal{N}(0, \widehat{\Xi}_n)$. Then for any $C\in\mathcal{C}$, we have

Thus

Put everything together.

Combining (ref), (ref), (ref), we have

Asymptotic regularity of $B_n$

Note that $B_n$ is block upper triangular, where the block sizes are corresponding to the decomposition of $\theta=( \theta_1, \dots, \theta_L)$. Recall that by construction, $\hat{f}\in[\sigma^2, M^2]$.

For any $j\leq j'$ we have:

Define

For $j\in\{1,\dots,L\}$, we have:

Next we shall show that for $j\leq j'$, we have

and that the event $\mathcal{E}^M_n$ happens with high probability:

Asymptotic negligibility of $\delta_{n,j,j'}$.

With $j\leq j'$, we have

where (i) uses the fact that $\delta_{n,j,j'}$ is a sum of martingale difference sequence, (ii) uses Lemmas (ref) & (ref), (iii) uses that $\Phi_{i,j'}$ is bounded and Lemma (ref), and (iv) uses Property (ref)(c).

High probability event $\mathcal{E}^M_n$.

We finally show that $\mathbb{P}(\left(\mathcal{E}^M_n\right)^c)=O(Ln^{\alpha_1-1})$. When $\mathcal{E}^M_n$ does not happen, by Lemma (ref), there exists a $j_0\in\{0,1,,\dots, l\}$ and $x$ with $\|x\|_2=1$ such that

Since $\delta_{n,j_0,j_0}=B_{n,j_0,j_0}-B_{n,j_0,j_0}^0$, by Lemma (ref), we have the following:

Rearranging yields the following results:

Thus by a union bound, we have

Proof of Corollaries

Proof of Corollary (ref)

Identity weights.

Consider $H_{i,j}$ being the identity matrix, we verify Property (ref).

Property (ref)(a). We have

Thus: \[ \left\{\frac{1}{n}\sum_{i=1}^n \mathbb{E}_{i,j}\left[H_{i,j} \mathrm{Var}^+_{i,j}(\Phi_{i,j}) \right]\right\}^\top \left\{\frac{1}{n}\sum_{i=1}^n \mathbb{E}_{i,j}\left[H_{i,j} \mathrm{Var}^+_{i,j}(\Phi_{i,j}) \right]\right\}\gtrsim n^{-2\alpha} I. \]

Property (ref)(b). We have

Property (ref)(c). We have

Consistency weights.

Consider $H_{i,j} = (I_{d_\tau}\otimes X_{i,j})\mathrm{Var}^+_{i,j}(T_{i,j})^{-1/2} (I_{d_\tau}\otimes X_{i,j})^\top $.

Property (ref)(a). We have

Thus: \[ \left\{\frac{1}{n}\sum_{i=1}^n \mathbb{E}_{i,j}\left[H_{i,j} \mathrm{Var}^+_{i,j}(\Phi_{i,j}) \right]\right\}^\top \left\{\frac{1}{n}\sum_{i=1}^n \mathbb{E}_{i,j}\left[H_{i,j} \mathrm{Var}^+_{i,j}(\Phi_{i,j}) \right]\right\}\gtrsim n^{-\alpha} I. \]

Property (ref)(b). We have

Property (ref)(c). We have

Proof of Corollary (ref)

Let $a$ be a given confidence level. Define set \[ \mathcal{A}(\ell) =\{x\in\mathbb{R}^{1+dL}: |\ell^\top x|\leq q_{\ell,\frac{a}{2}}\}, \]. By definition of $q_{\ell,\frac{\alpha}{2}}$, we have

The Strong Gaussian Approximation results says:

We thus have:

Proof of Corollary (ref)

Let $a$ be a given confidence level. Define set \[ \mathcal{A}(a) = \left\{ x\in\mathbb{R}^{1+dL}: \left\|(\hat{D}_n^{\dagger})^{1/2} x\right\|_\infty\leq q_{\infty,\, 1-a} \right\}. \] Then we have

On the other hand, we have:

Moreover, by Strong Gaussian Approximation, we have

Collectively, we have

Proof of Corollary (ref)

Recall that $V_{i,j}:=\mathbb{E}_{i,j}[X_{i,j}X_{i,j}^\top]$. Consider weighting scheme $H_{i,j} := \hat{f}_{i,j}^{-1/2}W_{i,j}$, where $\hat{f}_{i,j}$ satisfies Condition (ref) and \[ W_{i,j}=(I_{d_{\tau}}\otimes\hat{V}_{i,j}^{-1/2}) (I_{d_\tau}\otimes X_{i,j})\mathrm{Var}^+_{i,j}(T_{i,j})^{-1/2} (I_{d_\tau}\otimes X_{i,j})^\top\cdot\|X_{i,j}\|_2^{-2}, \] for $ \hat{V}_{i,j}$ adapted to $\mathcal{F}_{i,j}$ and satisfying $c^{-1}\cdot I \succeq\hat{V}_{i,j}\succeq c\cdot I$ for some $c>0$ and \[ \frac{1}{n} \sum_{i=1}^n \mathbb{E}\left[\left\|\hat{V}_{i,j}-V_{i,j} \right\|_{2}\right] = O(n^{-\gamma}). \] We verify $W_{i,j}$ satisfies Property (ref) with $\alpha_1=\alpha$ and $\alpha_3=\gamma$.

\paragraph{Checking Property (ref)(a).} We have

Since $\hat{V}_{i,j} \preceq c^{-1}\cdot I$, we have both $(I_{d_{\tau}}\otimes\hat{V}_{i,j}^{-1/2})$ and $\mathbb{E}_{i,j}\left[ (I_{d_\tau}\otimes X_{i,j}) \mathrm{Var}^+_{i,j}(T_{i,j})^{1/2} (I_{d_\tau}\otimes X_{i,j})^\top \right]$ being symmetric positive definite. Therefore $\frac{1}{n}\sum_{i=1}^n \mathbb{E}_{i,j}\left[W_{i,j} \mathrm{Var}^+_{i,j}(\Phi_{i,j}) \right]$ has eigenvalues being real number and positive. We have:

\paragraph{Checking Property (ref)(b).}

\paragraph{Checking Property (ref)(c).} We have

With $c\cdot I \preceq \hat{V}_{i,j} \preceq c^{-1}I$, $c\cdot I \preceq V_{i,j} \preceq c^{-1}I$ and $\| X_{i,j}\|_2^2\leq c^{-1}$, we have

and

Finally, we have

which also leads to $\frac{1}{n}\sum_{i=1}^n \|\mathbb{E}[W_{i,j} \mathrm{Var}_{i,j}^+(\Phi_{i,j}) W_{i,j}^\top] - I\|_2= O(n^{-\gamma})$ by Jensen's inequality. Therefore,