Automatic Doubly Robust Forests

Automatic Doubly Robust Forests

Introduction

A wide range of problems in social and life sciences can be framed as estimating the conditional expectation of a moment functional that depends on an unknown nuisance regression function:

Here, $m(z;g)$ is a known moment functional, and $g_0(x,w)=\mathbb{E}[Y\mid W=w,X=x]$ is an unknown nuisance regression function. The observation $Z=(Y,X,W)$ consists of an outcome variable $Y$, target covariates $X$, and additional variables $W$. A key example is heterogeneous treatment effect (HTE) estimation in causal inference, where the goal is to understand how a treatment $D$ affects the outcome $Y$ across different sub-populations or covariate groups defined by $X$. In HTE estimation, $W=(D,\tilde{X})$ includes the treatment variable $D$ and additional confounders $\tilde{X}$ that need to be controlled for. HTE estimation has crucial applications across various fields. For example, in personalized medicine, it helps identify effective treatments tailored to individual patient characteristics. In marketing, it guides customer segmentation and enables targeted interventions.

A primary challenge in estimating (ref) arises from the high dimensionality of the nuisance regression function. In many applications, it is desirable for $g_0$ to depend on high-dimensional variables $W$. For example, in HTE estimation, controlling for a large number of confounders $\tilde{X}$ is crucial for the validity of the causal identification argument. To achieve high-quality predictions in high-dimensional settings, modern machine learning (ML) methods with regularization and model selection are frequently used. However, these methods often introduce non-negligible bias due to the bias-variance trade-off inherent in optimizing mean squared error. This bias poses a significant challenge for inference, as it can lead to poor coverage of confidence intervals for the target estimand chernozhukov2018double,chernozhukov2022locally.

Much of the prior literature on non-parametric estimation of $\theta_0(x)$ assumes low-dimensional nuisance parameterizations wager2018estimation,grf. While this simplifies the estimation process, it limits applicability in real-world applications where a large number of additional variables $W$ are present. On the other hand, to accommodate high-dimensional nuisance functions, some existing works construct Neyman orthogonal moment functionals, which eliminate the first-order effect of nuisance estimation on the expected moment functional. However, these works typically impose parametric assumptions on the target function $\theta_0(x)$. For example, victor17 and chernozhukov2022locally focus on estimating the average treatment effect (ATE), where $\theta_0(x)$ is assumed to be constant across $x$. chernozhukov2017orthogonal and chernozhukov2018plug consider cases where $\theta_0(x)$ follows a linear specification, $\theta_0(x)=\theta_0^\top \phi(x)$, with $\theta_0$ assumed to be sparse and $\phi(x)$ a known feature mapping. Estimating a non-parametric $\theta_0(x)$ would offer more flexibility but is inherently more challenging than its parametric counterpart.

In this paper, we propose the Doubly Robust Random Forest (DRRF) algorithm to address these challenges of estimating conditional moment functionals. Our method extends the automatic debiased machine learning (Auto-DML) framework, which targets the conditional expectation of average moment functionals, to the more general setting of conditional moment functionals. To achieve this, we generalize the notion of the Riesz representer to a conditional Riesz representer and construct a doubly robust moment by incorporating a debiasing term based on this conditional representer. This debiasing mitigates the regularization bias inherent in plug-in machine learning methods and thereby accommodates high-dimensional nuisance functions. The estimation of the conditional Riesz representer is automatic, in the sense that it does not require prior knowledge of its form, and it can align with the estimation of the regression function \( g_0 \). We show that, under suitable conditions, the conditional Riesz representer can be projected onto the same functional space as \( g_0 \). For example, if \( g_0 \) is locally linear, we may, without loss of generality, restrict the conditional Riesz representer to be locally linear as well. We provide an efficient algorithm for estimating the conditional Riesz representer and constructing the debiasing term based on the forest lasso method of orf.

Our method is a two-stage procedure, similar to the Orthogonal Random Forest (ORF) algorithm proposed by orf. In the first stage, we estimate the nuisance functions and, consequently, the conditional moments. In the second stage, we compute similarity weights between each data point and the query point \( x \). The final estimator is constructed as a weighted average of the estimated conditional moments using a data-splitting approach. We estimate the weights following the strategy of orf, which enables nonparametric, forest-based estimation of the conditional functional \( \theta_0(x) \). However, our method significantly deviates from ORF in the manner it estimates the nuisance functions, making DRRF much more practical for large scale deployment when predicting $\theta_0(x)$ at a large number of query points $x_1,\ldots,x_m$. Specifically, the ORF algorithm requires re-estimating the ML-based nuisances for every single query point $x_i$, each time using different sample weights on the training data that are tailored to the query point. However, our algorithm estimates the nuisance functions only once for each of the training data points and does not re-estimate them for each query point after the training phase, as required by ORF. Since the estimation of nuisances is typically the most computationally intensive step, our approach substantially reduces computation when handling a large number of queries. This alternative nuisance estimation method also provides side benefits with respect to interpretability, which is important for most downstream decision making applications of the heterogeneous model $\theta_0(x)$. The final model $\hat\theta(\cdot)$ of the DRRF method is equivalent in representation to a standard regression forest, where the leaf nodes of each tree contain local estimates and the prediction of each tree at any query point $x$ is the value of the leaf node in which $x$ falls, and the prediction of the forest is simply the average of the tree predictions. Hence, standard interpretability tools for regression forests, such as SHAP lundberg2017unified, can be used off-the-shelf to interpret the model and provide insights into which features are important, on aggregate, and how each feature contributes to each query result $\theta_0(x_i)$. This is not the case for the ORF method, since evaluating $\theta_0(x)$ for any target $x$ involves complex nuisance re-estimation and is not a simple “look-up and average” calculation.

We show that our estimator is consistent and asymptotically normal, which allows for the construction of asymptotically valid confidence intervals. By incorporating the debiasing term, DRRF derives a doubly robust moment functional that relaxes the convergence rate requirements for the nuisance regression functions needed for consistency. Under appropriate regularity conditions, we show that our estimator achieves the same convergence rate as an oracle estimator with known nuisance functions, provided that the product of the estimation errors for the nuisance regression function and the conditional Riesz representer remains small. This property makes DRRF particularly advantageous in settings with high-dimensional and complex nuisance functions where accurate estimation of the conditional Riesz representer is feasible. Under additional conditions, we show that our estimator is asymptotically normal, enabling the construction of confidence intervals using bootstrap methods.

Variants of our DRRF estimation algorithm have already been widely used in practice. For instance, the ForestDRLearner method, implemented in the widely adopted open-source package EconML econml, is essentially the non-automatic analogue of our DRRF estimator for the HTE problem. However, no prior work has established the asymptotic normality of this method. The closest related prior work is the doubly robust variant of the ORF method, which however constructs the nuisance functions in a different manner, as described above—a distinction crucial to their proof. Hence, our asymptotic normality theorem not only provides a formal justification of these widely used methods but also enhances and generalizes them by incorporating automatic debiasing techniques. Automatic debiasing has also been shown to be empirically less sensitive to extreme propensity problems riesznet,RieszSharma, as it directly targets learning the Riesz function $\alpha_0$ rather than first estimating propensity models and then constructing the Riesz function through explicit division.

Our work contributes to two streams of recent literature at the intersection of causal inference and machine learning. First, our work builds on the literature on the estimation and inference of treatment effects that are robust to errors introduced by sophisticated machine learning methods. victor17 propose the double machine learning (double ML) method, which constructs a Neyman orthogonal moment by residualizing the effects of high-dimensional covariates on both treatments and outcomes. The treatment effect is then estimated through a linear regression between the residualized variables. chernozhukov2022automatic introduce the automatic debiased machine learning (Auto-DML) framework based on the (unconditional) Riesz representer of the score function, which ensures Neyman orthogonality and achieves double robustness for the target moment functional. riesznet further develop automatic procedures to learn the Riesz representation using neural networks and random forests. Building on this line of work, we extend the Auto-DML framework to conditional setting. We introduce the conditional Riesz representer and use it to construct debiased estimators for conditional functionals. Our conditional Riesz represneter can also be automatic estimated, resulting in a fully automatic debiasing procedure for conditional targets.

Second, our work contributes to the literature on machine learning-based estimation of conditional moment functions. grf introduce Generalized Random Forest (GRF), a flexible tree-based method for solving general conditional moment equations. GRF builds trees with partitions designed to capture heterogeneity in the target quantity. However, it cannot directly perform high-dimensional local nuisance estimations, so in practice, it is often combined with global nuisance function estimators in high-dimensional settings, which may not align well with the conditional estimand. orf extend GRF to the Orthogonal Random Forest (ORF) algorithm by combining the forest-based algorithm with Neyman orthogonal moments. While ORF allows for local fitting of high-dimensional nuisance functions around the target feature, it has several limitations. First, ORF employs double ML to construct Neyman orthogonal moments. Effective residualization in double ML typically requires well-specified models for the data generating process. For example, when estimating HTE, accurately modeling both the outcome and treatment variables is often crucial for achieving strong performance. Second, ORF refits nuisance functions for each target prediction point $x$, leading to high computational costs when the number of query points is large. In contrast, our DRRF method automatically constructs a doubly robust moment functional for estimating $\theta_0(x)$ without relying on any parametric or semi-parametric assumptions about the treatment variable. DRRF also facilitates efficient estimation of $\theta_0(x)$ across multiple query points. Additionally, compared to other doubly robust estimators that rely on the inverse of propensity score estimates, such as ForestDRLearner and the one proposed by zimmert2019nonparametric, our approach directly estimates the debiasing term, ensuring stability even when propensity scores are near boundary values.

We conduct extensive simulations to demonstrate the superior performance of the DRRF algorithm in HTE estimation over benchmark approaches, including variants of ORF and GRF. Our results highlight three main advantages of DRRF. First, like ORF, DRRF performs local fitting of nuisance functions and generally outperforms GRF variants that rely on global fitting. Second, DRRF performs automatic debiasing and does not require a well-specified model for the treatment variable. In contrast, ORF may suffer from significant bias or variance if the treatment model is misspecified, whereas DRRF remains stable and provides accurate estimates. Lastly, DRRF is significantly more computationally efficient than ORF, with much shorter evaluation runtime when estimating at multiple target query points.

The rest of the paper is organized as follows. In Section (ref), we introduce our doubly robust conditional moment estimation procedure. In particular, we define the general framework in Section (ref) and propose the Doubly Robust Random Forest algorithm in Section (ref) as an efficient implementation of the general framework. In Section (ref), we provide theoretical results that establish the consistency and asymptotic normality of our estimator under appropriate conditions. Finally, in Section (ref), we conduct comprehensive synthetic data simulations to demonstrate the advantages of our method over benchmark approaches.

Problem Statement

Consider a dataset with $2n$ i.i.d. observations $Z_i=(Y_i,X_i,W_i)$, where each observation $Z_i$ contains a scalar outcome of interest $Y_i$, target covariates $X_i$, and additional variables $W_i$. This paper provides an automatic, doubly robust, and computationally efficient approach for estimating the conditional expectation of a moment functional

evaluated at multiple target points $x=x_1,\ldots,x_m$. Here, $g_0(x,w)$ is an unknown nuisance regression function defined by $g_0(x,w) \coloneqq \mathbb{E}[Y\mid W=w,X=x]$, and $m(z;g)$ is a known moment functional that depends on a data observation $z$ and a possible regression function $g$. We assume that the target feature space $\mathcal{X}=[0,1]^d$ has a fixed dimension $d$, and that $X$ has a density bounded away from both zero and infinity. We allow the dimension of $W$ to be large and grow with $n$, positioning the problem in a high-dimensional setting where the nuisance regression function can depend on a significant number of variables.

To illustrate the general target in (ref), we consider estimation problems related to heterogeneous treatment effects (HTE) in Examples (ref)-(ref). These examples demonstrate the broad applicability of our framework to moment functionals in causal inference. They also highlight the importance of including a large number of covariates in the nuisance regression function $g_0$, as doing so enhances the plausibility of the unconfoundedness assumption for treatment effect identification.

\paragraph{Notations.} In the following, we define $\left\lVertg\right\rVert_x = (\mathbb{E}[g(x,W)^2\mid X=x])^{1/2}$ and let $\mathcal{G}_x$ be the set of functions such that $\left\lVertg\right\rVert_x<\infty$. We use $f(n)=\omega(g(n))$ to denote that $f$ asymptotically dominates $g$, i.e., $f(n)/g(n)\rightarrow \infty$ as $n\rightarrow \infty$, and $f(n)=\Theta(g(n))$ denotes that $f$ is asymptotically bounded by $g$ both from above and below, i.e., $f(n)=O(g(n))$ and $g(n)=O(f(n))$.

Automatic Doubly Robust Conditional Moment Estimation

Conditional Riesz Representer and Debiased Moment Functional

In this section, we show how to adapt the automatic debiasing method using the Riesz representer chernozhukov2022automatic to estimate the conditional target $\theta_0(x)$. To apply the Riesz representation theorem, we assume that, for any $x \in \mathcal{X}$, the conditional expected moment $\mathbb{E}[m(Z;g(x,W))\mid X=x]$ is a continuous linear functional of $g$ with domain $\mathcal{G}_x$. Note that the conditional function space $\mathcal{G}_x$ is the $L_2$ space under the law $\mathcal{L}([X,W]\mid X=x)$, and specifically, it is a Hilbert space with inner product defined by $$\langle g_1,g_2\rangle_x = \mathbb{E}[g_1(X,W)g_2(X,W)|X=x]. $$

By the Riesz representation theorem, there exists a unique function $\alpha_0\in \mathcal{G}_x$, such that for any $g\in \mathcal{G}_x$,

We refer to $\alpha_0$ as the conditional Riesz representer (cRR) to highlight its distinction from the standard Riesz representer used for global moment functionals: the cRR depends explicitly on the conditioning variable $X = x$. To first build some intuition for the cRR $\alpha_0$ and its role in our framework, we illustrate it in the context of CATE estimation in Example (ref).

Using the conditional Riesz representer, we construct a debiased moment functional $\psi(Z; \alpha_0, g)$ for the original moment functional $m(Z; g)$. This construction is analogous to the use of unconditional Riesz representer in the debiased machine learning literature for estimating unconditional estimands. Specifically, we define \[ \psi(Z; \alpha_0, g) \coloneqq m(Z; g(X, W)) + \alpha_0(X, W)(Y - g(X, W)). \]

It is straightforward to show that the debiased moment functional $\psi$ satisfies $\mathbb{E}[\psi(Z;\alpha_0,g_0)\mid X=x]=\theta_0(x)$ by conditioning on $X$ and $W$ and applying the law of iterated expectations. Therefore, we can estimate $\theta_0(x)$ through a plug-in weighted average of $\psi(Z;\hat{\alpha},\hat{g})$, using estimators for both $g_0$ and $\alpha_0$. Importantly, the debiased moment functional $\psi(Z;\alpha,g)$ is locally Neyman orthogonal, meaning that the regression estimator $\hat{g}$ and the cRR estimator $\hat{\alpha}$ have no first-order effect on the moment at the target point $x$. Furthermore, $\psi(Z;\alpha,g)$ is doubly robust for the estimation of $\theta_0(x)$, which allows our estimator to achieve a fast convergence rate under mild consistency requirement for the nuisance estimators $\hat{g}$ and $\hat{\alpha}$. We will elaborate on this property in our theoretical results in Section (ref). In Propositions (ref) and (ref) below, we formally state the local Neyman orthogonality and double robustness properties of the debiased moment functional $\psi(Z;{\alpha},{g})$.

In Proposition (ref), the local orthogonality implies that the conditional expectation of the debiased moment functional is insensitive to local estimation errors in the nuisance functions $\hat g$ and $\hat{\alpha}$ around their true values. The double robustness in Proposition (ref) ensures the bias of the plug-in estimators of $\psi$ converges to zero at the rate determined by the product of the convergence rates of $\hat\alpha$ and $\hat g$. This property is crucial since it indicates that a fast convergence rate for the estimation of $\theta_0$ can be achieved without requiring highly accurate estimates of either $g_0$ or $\alpha_0$ individually. The double robustness property accommodates complex nuisance regression functions with high-dimensional parameterizations and is particularly desirable in scenarios where accurate estimation of the conditional Riesz representer is achievable.

This debiasing approach is automatic in the sense that, it does not require prior knowledge of the functional form of the debiasing term. In particular, the conditional Riesz representer $\alpha_0$ can be estimated directly as the minimizer of the conditional loss function

In Section (ref), we show that, under appropriate conditions, the estimation of the cRR $\alpha_0$ is naturally aligned with the estimation of the regression function $g_0$ based on the Hilbert projection theorem. Specifically, when $g_0$ follows a partially linear form, $\alpha_0$ can, without loss of generality, be restricted to the same functional form. We leverage this insight to develop an algorithm that automatically estimates the conditional Riesz representer and, in turn, constructs the debiasing term for the original moment functional.

Doubly Robust Estimation

Now, we introduce our doubly robust procedure for estimating $\theta_0(x)$. We first outline the procedure for estimating $\theta_0(x)$ at a single target point $x$, and then demonstrate how to efficiently estimate $\theta_0(x)$ at a large number of query points.

To estimate $\theta_0(x)$, we begin by splitting the $2n$ observations $Z_i = (Y_i,X_i,W_i)$ into two halves. The second half of the data is used to construct the regression estimator $\hat g$ and the cRR estimator $\hat\alpha$, while the first half is used to compute similarity weights and construct a plug-in estimator of $\theta_0(x)=\mathbb{E}[\psi(Z;\alpha_0,g_0)|X=x]$. Specifically, the procedure consists of the following two steps:

In the second step, the weight $K(x,X_i,Z_{1:n},\xi)$ is estimated based on $Z_{1:n}$ and measures the importance of the sample point $X_i$ for estimating $\theta_0(x)$. In this paper, we consider the sub-sampled kernel that is obtained by averaging $B$ sub-samples of the base kernel $\tilde K$:

where $\{Z_{S_b}\}_{b=1}^{B}$ are random size-$s$ sub-samples of $\{Z_1,\ldots,Z_n\}$, $\xi_b$ represents the internal randomness of the base kernel, and $\xi$ captures the randomness from both sub-sampling and $\{\xi_b\}_{b=1}^{B}$. The base kernel $\tilde K$ satisfies $\tilde K(x,X_i,Z_{S_b},\xi_b)=0$ for any $i\notin S_b$ and $\sum_{i\in S_b} \tilde K(x,X_i,Z_{S_b},\xi_b) = 1$ for all $b$.

In Section (ref), we propose a data-efficient tree-based algorithm to construct the base kernel $\tilde K$, which directly targets maximizing the heterogeneity of $\theta_0(x)$. Additionally, we demonstrate in Section (ref) the use of the forest lasso algorithm for estimating the nuisance functions in the first step of our two-step procedure.

\paragraph{Estimation at multiple points.} In many practical applications, we are interested in estimating $\theta_0(x)$ at a large number of points $x=x_1,\ldots,x_m$. Observe that our estimator in (ref) depends on $x$ only through the kernels \{$K(x,X_i,Z_{1:n},\xi)\}_{i=1}^n$, which measure the similarities of $x$ and the $X_i$'s. Therefore, we can pre-compute and store $\{ m(Z_i;\hat{g}(X_i,W_i)) + \hat{\alpha}(X_i,W_i)(Y_i-\hat{g}(X_i,W_i))\}_{i=1}^n$ based on the training data. To estimate $\theta_0(x_1),\ldots \theta_0(x_m)$, it is then sufficient to evaluate the kernel $K(x,X_i,Z_{1:n},\xi)$ at $x=x_1,\ldots,x_m$ for each $i\in \{1,\ldots,n\}$. Unlike the Orthogonal Random Forest (ORF) orf, which requires the estimation of the nuisance functions at each new target point $x$, our procedure allows for the pre-computation and storage of the nuisance function estimates at the sample covariates $X_i$. As a result, when querying multiple points $x_1, \ldots, x_m$, our method avoids re-estimating the nuisance functions, significantly accelerating the estimation process. This innovation, however, makes the derivation of the asymptotic normality results for our estimator more challenging, as it requires establishing a uniform bound on nuisance errors. This is accomplished in our proof of Theorem (ref).

Doubly Robust Random Forest

In this section, we introduce the Doubly Robust Random Forest (DRRF) for estimating $\theta_0(x)$, which provides an efficient implementation of the doubly robust estimation procedure in Section (ref). Specifically, we outline the estimation procedures for the nuisance functions $g_0$ and $\alpha_0$ in Section (ref) and detail the construction of the forest kernel in Section (ref).

Estimation of \texorpdfstring{$g_0$}{g0} and \texorpdfstring{$\alpha_0$}{alpha0}

In practice, our framework allows for the use of a variety of methods for estimating the nuisance functions $g_0$ and $\alpha_0$. However, to provide a concrete example and establish the foundation for the theoretical results of our estimator, we adopt a specific form of the nuisance regression $g_0$ following orf. Specifically, we assume that $g_0(x,w) = h(w,\nu_g(x))$ for some known function $h$ and unknown function $\nu_g:\mathcal{X}\rightarrow \mathbb{R}^{d_{\nu}}$. We consider the function $h$ such that $\mathcal{H}_x \coloneqq \{g\in \mathcal{G}_x: g(x,w)=h(w,\nu(x)) \ \text{for some function} \ \nu:\mathcal{X}\rightarrow \mathbb{R}^{d_{\nu}} \}$ is a closed linear subspace of $\mathcal{G}_x$. For example, this holds when $h(w, \nu(x)) = r(w)^\top \nu(x)$ for some function $r$, under mild conditions.\footnote{See Appendix (ref) for a proof.} We allow the dimension $d_{\nu}$ to grow with $n$.

Consequently, the conditional Riesz representer $\alpha_0(x, w)$ can be projected onto the space $\mathcal{H}_x$, yielding $\alpha_{\mathcal{H}_x}(x,w)$ such that for all \( g \in \mathcal{H}_x \), \[ \mathbb{E}[m(Z; g(X, W)) \mid X = x] = \mathbb{E}[g(X, W) \alpha_{\mathcal{H}_x}(X, W) \mid X = x]. \]

Therefore, without loss of generality, we can consider the conditional Riesz representer $\alpha_0(w,x)$ as taking the form $h(w,\nu_\alpha(x))$, with the same function $h$ and an unknown function $\nu_\alpha:\mathcal{X}\rightarrow \mathbb{R}^{d_{\nu}}$.

Suppose that the unknown functions $\nu_g$ and $\nu_\alpha$ can be identified as the minimizers of the following local losses with known loss functions $l_g$ and $l_\alpha$:

where $V$ is a bounded space.

Next, we demonstrate the use of the Forest Lasso algorithm proposed by orf to estimate $g_0$ and $\alpha_0$ in the first step of the doubly robust estimation procedure. Specifically, we run the tree learner defined in Section (ref) on $\{Z_{n+1},\ldots, Z_{2n}\}$ based on $B$ subsamples of size $s$ to obtain kernels $K(\cdot,X_i,Z_{n+1:2n},\xi)$ for $i = n+1, \ldots, 2n$. Then, we solve the lasso problems

at $X_1,\ldots,X_n$. In contrast to orf, which applies Forest Lasso at the target point $x$, our approach runs the algorithm at the training data points $X_1,\ldots,X_n$. This modification eliminates the need to rerun the algorithm for each new target point once we solve the $n$ Forest Lasso programs associated with $X_1,\ldots,X_n$.

As to be shown in Section (ref), with appropriately chosen regularization parameters $\lambda_g$ and $\lambda_{\alpha}$, the estimators $\hat{\nu}_g$ and $\hat{\nu}_\alpha$ are consistent under the assumptions that the high-dimensional functions $\nu_g$ and $\nu_\alpha$ are sparse and supported on a finite set of dimensions, along with other regularity conditions. These results are formalized in Lemma (ref) in Appendix, which will then be used to establish the asymptotic normality of the target parameter $\theta_0(x)$ in Theorem (ref).

Forest kernel

We proceed to describe the forest algorithm for computing the similarity weights $K(x,X_i,Z_{1:n},\xi)$ in (ref) and in the second step of our estimation. As stated in Section (ref), the kernel $K(x,X_i,Z_{1:n},\xi)$ is the average over $B$ base kernels $\tilde K(x,X_i,Z_{S_b},\xi_b)$ constructed from random size-$s$ sub-samples $S_b$ for $b=1,\ldots,B$. Our algorithm combines the debiased moment functional with the tree-based learner from ORF orf to compute the base kernel $\tilde K$.

For each sub-sample $S_b$, we randomly partition it into two equal-sized subsets $S_b^1$ and $S_b^2$. We use $S_b^1$ to grow a tree learner that splits the feature space $\mathcal{X}$ into small subspaces, and use $S_b^2$ to evaluate the similarity weights.

\paragraph{Splitting algorithm.} For $S_b^1$, we start from a root node containing the entire space $\mathcal{X}$. At each node $P$, we implement the following three-stage procedure:

Throughout the tree-growing process, we maintain the tree properties specified in Assumption (ref).

\paragraph{Weight estimation.} For each tree based on $S_b^1,$ let $L_b(x)$ denote the leaf that contains the target value $x$. We compute the base kernels $\tilde K(x,X_i,Z_{S_b},\xi_b)$ for each point $i$ as

Note that we only assign positive weights $\tilde K(x,X_i,Z_{S_b},\xi_b)$ to the sample points in $S_b^2$.

Theoretical Results

In this section, we establish the consistency and asymptotic normality of the general doubly robust estimator introduced in Section (ref), as well as the DRRF estimator described in Section (ref), which provides an effective implementation of the general estimator. These results are derived under appropriate regularity conditions.

Consistency

We begin by establishing the consistency of the doubly robust estimator $\hat{\theta}(x)$, constructed using sub-sampled kernels as described in Section (ref), for estimating the target parameter $\theta_0(x)$. To ensure consistency, we require that the base kernel $\tilde K$ satisfies the following properties:

Assumption (ref) ensures the conditional independence of the base kernel and the estimated debiased moment functional. Assumption (ref) states that, in expectation, positive weights are assigned only to sample points whose features $X_i$ are sufficiently close to the target $x$. It is straightforward to verify that Assumption (ref) holds for the DRRF algorithm defined in Section (ref). Moreover, following wager2018estimation, the DRRFs that satisfy the forest regularity conditions specified in Assumption (ref) achieve a kernel shrinkage rate of $\epsilon(s) = O(s^{-1/(2a d)})$,\footnote{The parameter $a=\log(\rho^{-1})/(\pi\log((1-\rho)^{-1}))\,,$ where $\rho$ represents the smallest fraction of leaves on each side of the split, and $\pi/d$ is the smallest probability that the split occurs along each feature, both as defined in Assumption (ref).} which decays polynomially in $s$.

Finally, we impose additional conditions on the debiased moment functional, as formalized in Assumption (ref). These conditions introduce standard boundedness and smoothness assumptions, similar to those assumed in, for example, orf.

Now we formally state the consistency result for the general doubly robust estimator $\hat{\theta}(x)$ proposed in Section (ref).

In Theorem (ref), the right-hand side of inequality (ref) decomposes into four terms, each reflecting a distinct source of error. The first term captures the variation introduced by aggregating over the $B$ sub-samples used to construct the base kernels. The second term represents the sampling error associated with an appropriate complete U-statistic based on the observed data. The third term reflects the error coming from kernel estimation, where $\epsilon(s)$ denotes the kernel shrinkage rate specified in Assumption (ref). Finally, the fourth term accounts for the estimation errors in the regression function and the conditional Riesz representer. Note that the convergence rate of $\hat{\theta}(x)$ depends on the nuisance functions solely through the product $\mathcal{E}(\hat{g},\hat{\alpha}) = \left\| \hat{g}-g_0\right\|_x \cdot \left\| \hat{\alpha}-\alpha_0\right\|_x.$ Therefore, the proposed estimation procedure is indeed doubly robust with respect to nuisance estimation, and it achieves the oracle rate as long as $\mathcal{E}(\hat{g},\hat{\alpha})$ converges to 0 at a slow rate.

We now establish the consistency of the doubly robust estimator $\hat{\theta}(x)$ estimated by our DRRF algorithm introduced in Section (ref). Following the discussion in Section (ref), we focus on nuisance estimators of the form $\hat{g}(x,w)=h(w,\hat{\nu}_g(x))$ and $\hat{\alpha}(x,w)=h(w,\hat{\nu}_\alpha(x))$. To provide theoretical guarantees for the local regularized nuisance estimators in (ref), we impose the following Assumption (ref) on the true parameters $\nu_g$ and $\nu_\alpha$.

At a high level, Assumption 4 ensures that the nuisance parameters $\nu_g$ and $\nu_\alpha$ are sparse, and can be reliably estimated from the data under standard identifiability conditions. This assumption is important for proving that the nuisance estimators $\hat{g}(x,w)=h(w,\hat{\nu}_g(x))$ and $\hat{\alpha}(x,w)=h(w,\hat{\nu}_\alpha(x))$, defined via the penalized regression problems in (ref) in Section (ref), converge to the true functions $g_0$ and $\alpha_0$, respectively. Under this assumption, we establish the consistency result for DRRF estimator.

Corollary (ref) shows that the DRRF algorithm achieves the same consistency rate as estimating the conditional functional with oracle nuisance functions. The convergence rate $\beta$ depends on the forest regularity parameter $a = \log(\rho^{-1}) / (\pi \log((1 - \rho)^{-1}))$ (see Assumption (ref)) and the dimension $d$ of the target feature space $\mathcal{X}$. Specifically, the rate improves with more balanced tree splits (i.e., larger $\rho$ and $\pi$) and lower dimension $d$, given suitable choices of the subsample size $s$ and the number of base kernels $B$.

Asymptotic Normality

In this section, we establish the asymptotic normality results for the general doubly robust estimator and DRRF. We again focus on nuisance estimators of the form $\hat{g}(x,w)=h(w,\hat{\nu}_g(x))$ and $\hat{\alpha}(x,w)=h(w,\hat{\nu}_\alpha(x))$. Additionally, we assume that the dimensionality $d_{\nu}$ of $\nu_g$ and $\nu_\alpha$ grows at most polynomially in $n$. We require an additional assumption on the debiased moment functional:

Assumption (ref) imposes a local Lipschitz-type condition on the moment functional $\psi(Z;\alpha,g)$ with respect to the nuisance functions $g$ and $\alpha$. Specifically, it ensures that the conditional squared error in the moment functional arising from estimation of $g$ and $\alpha$ can be bounded by the corresponding estimation errors of these functions. This assumption is imposed for the asymptotic results in Theorems (ref) and (ref), as it guarantees that the sampling error associated with estimated nuisances can be asymptotically dominated by the term involving the oracle nuisances $\alpha_0$ and $g_0$, whose asymptotic variance involves the term $$\eta(s)\coloneqq Var\left( \mathbb{E} \left[ \sum_{i=1}^s\tilde{K}(x,X_i,Z_S,\xi_S)\psi(Z_i;\alpha_0,g_0)\bigg|Z_1 \right]\right)\,,$$ where the set $S$ is fixed to include the first $s$ samples.

We now state the asymptotic normality result for the general doubly robust estimator proposed in Section (ref).

In Theorem (ref), the conditions $n\eta(s) \rightarrow \infty$ and $n\epsilon^2(s)/(s^2\eta(s))\rightarrow 0$ require $\eta(s)$ to decay slower than $1/n$ and $\epsilon^2(s)/(s^2\eta(s))$ to decay faster than $1/n$. This can be achieved under the “double-sample” scheme in the DRRF algorithm under an appropriate choice of the subsample size $s$. Moreover, the nuisance rate condition $\sqrt{n/(s^2\eta(s))}\cdot r_n^2(\delta_n/n)\rightarrow 0$ in Theorem (ref) is usually mild and only requires the nuisance parameters $\nu_g$ and $\nu_\alpha$ to converge at a slow rate. Following Theorem (ref), we immediately obtain the following asymptotic normality result for DRRF.

In light of Theorem (ref), we can use the Bootstrap of Little Bags sexton2009standard, grf approach to construct asymptotically valid confidence intervals for $\theta_0(x)$. In Section (ref), we show that this construction provides good empirical coverage.

Simulations

In this section, we conduct comprehensive simulations to evaluate the performance of our DRRF estimator in heterogeneous treatment effect (HTE) estimation, comparing it against other tree-based non-parametric methods.

We generate data according to the following process: \[ Y_i = \theta_0(X_i)\cdot D_i+ \tilde X_i^\top \nu_0 + \epsilon_i, \] where the target covariates $X_i \overset{i.i.d.}{\sim} U[0,1]$, the noise terms $\epsilon_i \overset{i.i.d.}{\sim} U[-1,1]$, and the coefficient vector $\nu_0$ is $k$-sparse, with non-zero coefficients drawn i.i.d. from $U[-1,1]$. Following orf, we define the true heterogeneous treatment effect $\theta_0(x)$ to be a piecewise linear function: \[ \theta_0(x) =

\] Our goal is to estimate $\theta_0(x)$ at various query points $x$. We consider different data-generating processes for the treatment variable $D_i$ and the control covariates $\tilde{X}_i$. In the main text, we focus on the case where the treatment variable $D_i$ is binary, while Appendix (ref) presents results for continuous treatment settings. The total number of observations is fixed at $2n = 5000$, with $p = 100$ additional covariates. The support size $k$ varies across $\{5, 10, 15\}$. A detailed description of the data-generating processes is provided in Table (ref).

We compare the performance of two variants of DRRF, along with two variants each of ORF orf and GRF grf:

There are two key differences among the three variants. First, both DRRF and ORF variants accommodate high-dimensional nuisance functions and perform local nuisance estimation. In contrast, GRF variants apply tree learner-generated similarity weights to moment functions constructed with global nuisance estimators. Second, while both DRRF and GRF variants estimate the target based on a doubly robust moment, ORF variants use an orthogonal moment without the doubly robust property. On the other hand, constructing the doubly robust moment in GRF variants requires propensity score estimation, which restricts their applicability to continuous treatment settings. In contrast, our DRRF variants automatically construct the doubly robust moment without requiring propensity score estimation, making them applicable to both discrete and continuous treatment settings.

Table (ref) compares the estimation accuracy and evaluation runtime of different estimators. Specifically, it reports the root mean squared error (RMSE) of the estimated $\hat{\theta}(x_i)$ and the average evaluation runtime across 100 test points $x_i$ over a uniform grid in $[0,1]$. We report results based on 100 simulations for each setup. Figures (ref)-(ref) illustrate the performance of each estimator by plotting the average estimated effects, along with the 5th and 95th percentiles of the estimated effects across 100 replications. Additionally, Figure (ref) provides an example of the bootstrap confidence intervals for the DRRF algorithm, demonstrating reliable finite-sample coverage.

In the following, we highlight three key advantages of DRRF variants. First, similar to ORF variants, DRRF variants perform local fitting of nuisance functions, which generally leads to superior performance compared to GRF variants, which rely on global fitting. The global nuisance estimation in GRF variants optimizes overall mean squared error, which does not align with the objective of final-stage local estimation. As reported in Table (ref), both DRRF variants exhibit significantly lower RMSE than GRF variants in the binary treatment setting. In fact, Figures (ref)-(ref) show that GRF variants suffer from substantially higher variance. This high variance may also be linked to their dependence on inverse propensity score estimates, which can become unstable when the propensity scores are near boundary values.

Second, DRRF variants do not require a well-fitted model for the treatment variable, making them more robust than ORF variants. As shown in Table (ref), DRRF variants outperform ORF variants in RMSE in most cases. Note that in the three setups considered, the relationship between treatment $D_i$ and covariates $\tilde{X}_i$ is non-linear. In such cases, the Lasso-based residualization in ORF variants struggles to adequately residualize $D$ from $\tilde{X}$, degrading the performance of ORF variants. This issue is more pronounced in the continuous treatment setting, as shown in Table (ref) in Appendix (ref). In contrast, our DRRF algorithms automatically construct a debiasing term without relying on a well-specified model between $D$ and $\tilde{X}$.

Finally, DRRF variants are significantly more computationally efficient than ORF variants, especially ORF-CV, when estimating $\theta_0(x)$ at multiple query points. As shown in Table (ref), the evaluation runtimes of both DRRF and DRRF-CV are substantially shorter than those of ORF and ORF-CV. In particular, DRRF-CV’s average evaluation runtime is over 100 times faster than ORF-CV. This demonstrates the efficiency of the DRRF algorithm and makes it a more practical choice for real-world applications requiring fast estimations at a large number of query points.

Collectively, these results demonstrate the advantages of the proposed DRRF algorithms in terms of estimation accuracy, robustness, and computational efficiency, underscoring their potential for practical applications.

Conclusion

In this paper, we proposed an efficient and automatic doubly robust method for estimating conditional moment functionals, introducing the Doubly Robust Random Forest (DRRF) as a practical algorithm. Our approach extended the existing automatic debiased machine learning (Auto-DML) framework, based on Riesz representers and originally developed for average functionals, to the more general setting of conditional functionals. In particular, we introduced the notion of conditional Riesz representer (cRR) and showed that, under mild conditions, cRR can be projected onto the same function space as the regression function. In particular, when the regression function is locally linear, it suffices to consider a locally linear conditional Riesz representer, which can be efficiently estimated via a weighted Lasso. Our method also offers substantial computational gains at evaluation time compared to existing approaches such as orf by precomputing and storing moment estimates from the training data. We established the consistency and asymptotic normality of our procedure under appropriate regularity conditions. Overall, our method provides a nonparametric, doubly robust, and computationally scalable solution for estimating general conditional moment functionals in the presence of high-dimensional nuisance functions.

\setcounter{section}{0} \setcounter{assumption}{0} \setcounter{lemma}{0} \setcounter{proposition}{0} \setcounter{table}{0} \setcounter{figure}{0}