Optimal Policy Adaptation Under Covariate Shift

1 Introduction

In many real-world scenarios, labeled data is often scarce due to budget constraints and time-consuming collection processes [Zhuang et al., 2020; Imbens et al., 2024], significantly limiting the generalizability of the resulting models. For example, in medical research, collecting labeled data involves extensive clinical trials and follow-up periods, making it costly and time-consuming [Dahabreh et al., 2020; Hu et al., 2023]. In autonomous driving, obtaining labeled data requires manual annotation of large amounts of sensor data, which is laborious and expensive [Sun et al., 2020]. To address this problem and enhance a model's performance in a target domain without labels, an active area of research is transfer learning. It aims to improve the performance of target learners in the target domain by transferring the knowledge contained in a different but related source domain.

∗ corresponding author

While transfer learning has been extensively studied in the context of prediction models [Wang et al., 2018; Wang et al., 2020; Pesciullesi et al., 2020], how to transfer a policy is still underdeveloped. Policy learning refers to identifying individuals who should receive treatment/intervention based on their characteristics by maximizing rewards [Murphy, 2003]. It has broad applications in recommender systems [Chen and Sun, 2021; Wu et al., 2022], precision medicine [Bertsimas et al., 2017] and reinforcement learning [Liu et al., 2021; Kwan et al., 2023]. Unlike transfer learning for prediction models, policy transfer faces identification challenges due to its counterfactual nature [Athey and Wager, 2021; Li et al., 2023b; Wu et al., 2024c; Yang et al., 2024]. Instead of predicting outcomes based on observed data, policy transfer requires considering what would happen under different actions, making the process more complex.

We aim to learn the optimal policies in the target and entire domains using a dataset from the source domain (source dataset) and a dataset from the target domain (target dataset). The source dataset includes the covariates, treatment, and outcome for each individual, whereas the target dataset contains only the covariates. We assume that the source dataset satisfies the unconfoundedness and overlap assumptions while imposing fewer restrictions on the target dataset. We allow for substantial differences in the covariate distributions between the source and target datasets (referred to as covariate shift), while assuming that the conditional distributions of potential outcomes given covariates are the same.

In this article, we first propose a principled policy learning approach under covariate shift. Specifically, we define the reward and the optimal policy in the target domain using the potential outcome framework in causal inference. Under the widely used assumptions of unconfoundedness and transportability, we establish the identifiability of the reward in the target domain and then derive its efficient influence function and semiparametric efficiency bound. Building on this, we develop a novel estimator for the reward. Theoretical analysis shows that the proposed estimator is doubly robust and achieves the semiparametric efficient bound, that is, it is the optimal regular estimator in terms of asymptotic variance

[Newey, 1990]. Then we propose to learn the optimal policy by maximizing the estimated reward. We analyze the bias of the estimated reward and the generalization error bound of the learned policy. In addition, we extend the proposed method to learn the optimal policy in the entire domain consisting of the source and target domains by leveraging data from both domains to address distributional discrepancies and ensure robust generalization across heterogeneous environments.

The main contributions are summarized as follows: (1) We propose a principled approach for learning the optimal policy under covariate shift from a perspective of causality, by introducing plausible identifiability assumptions and efficient estimation methods; (2) We provide a comprehensive theoretical analysis of the proposed approach, including the consistency, asymptotic normality, and semiparametric efficiency of the estimator of reward. Additionally, we derive the bias and generalization error bound for the learned policy; (3) We conduct extensive experiments to demonstrate the effectiveness of the proposed policy learning approach.

2 ProblemFormulation

2.1 Notation and Setup

Let A ∈ A = {0, 1} denote the binary indicator for treatment, where A = 1 indicates receiving treatment and A = 0 indicates not receiving treatment. The random vector X ∈ X ⊂ Rp represents the p-dimensional covariates measured before treatment, and Y ∈ Y ⊂ R denotes the outcome of interest. Assume that a larger outcome is preferable. Under the potential outcome framework [Rubin, 1974; Splawa-Neyman, 1990], let Y (a) denote the potential outcome that would be observed if A were set to a for a ∈ A. By the consistency [Hern´an and Robins, 2020], the observed outcome Y satisfies Y = Y (A) = AY (1) + (1 - A)Y (0).

Without loss of generality, we consider a typical scenario involving two datasets: a source dataset and a target dataset, which are representative samples of the source domain and target domain, respectively. Let G ∈ {0, 1} be the indicator for the data source, where G = 1 denotes the source domain and G = 0 denotes the target domain. The observed data are represented as follows,

where the source dataset D1 consists of n1 individuals, with observed covariates, treatment, and outcome for each individual. The target dataset D0 contains n0 individuals, with only covariates for each individual. This is common in real life due to the scarcity of outcome data. For example, in medical research, patient features are observed, but obtaining outcomes requires long-term follow-up [Hu et al., 2023; Imbens et al., 2024]. Let P(·|G = 1) and P(·|G = 0) be the distributions of the two datasets respectively. Then n = n0 + n1 and q = n1/n represent the probability of an individual belonging to the source population.

2.2 Formulation

We formulate the goal of learning the optimal policy in the target domain. Specifically, let π : X → A denote a policy that maps individual covariates X = x to the treatment space A. A policy π(X) is a treatment rule that determines whether an individual receives treatment (A = 1) or not (A = 0). For a given policy π applied to the target domain, the average reward is defined as follows

We aim to learn the optimal policy π∗ defined by π∗ = argmaxπ∈ΠR(π), where Π is a pre-specified policy class. For example, π(X) can be modeled with a parameter θ using methods such as logistic regression or multilayer perceptron, with each value of θ corresponding to a different policy.

In addition, for a policy π(x) applied across the whole domain, the corresponding average reward is defined as

There is a subtle difference between R(π) and V (π). For R(π), our focus is on transferring the policy from the source domain to the target domain, and for V (π), we aim to generalize the policy from the source domain to the entire domain. In the main text, we focus on learning the policy maximizing R(π) to avoid redundancy. We also develop a similar approach to learn the policy maximizing V (π) and briefly present it in Section 4.3.

3 Oracle Policy and Identifiability

3.1 Oracle Policy

The optimal policy that maximizes Eq. (1) has an explicit form. Let τ(X) = E[Y (1) - Y (0)|X, G = 0] be the conditional average treatment effect (CATE) in the target domain,

where the last equality follows from the law of iterated expectations. Then we have the following conclusion.

Lemma 1. The oracle policy

where maxπ is taken over all possible policies without constraints, rather than being restricted to Π.

For an individual characterized by X = x in the target domain, Lemma 1 asserts that the decision to accept treatment (A = 1) should be based on the sign of τ(x). The oracle policy π∗ 0 recommends treatment for individuals expected to experience a positive benefit, thereby optimizing the overall reward within the target domain. The target policy π∗ equals the oracle policy π∗ 0 in Lemma 1 if π∗ 0 ∈ Π; otherwise, they may not be equal, and their difference is the systematic error induced by limited hypothesis space of Π.

3.2 Identifiability of the Reward

To learn the optimal policy π∗ , we first need to address the identifiability problem of R(π), as this forms the foundation for policy evaluation. Since the target dataset only contains covariates X, R(π) cannot be identified from the target data alone due to the absence of treatment and outcome. To identify R(π), it is necessary to borrow information from the source dataset by imposing several assumptions.

Assumption 1. For all X in the source domain,

Assumption 1(i) states that, in the source domain, the treatment is independent of the potential outcomes given the covariates, implying that all confounders affecting both the treatment and outcome are observed. Assumption 1(ii) asserts that any individual characterized by X in the source domain has a positive probability of receiving treatment. Assumption 1 is a standard assumption for identifying causal effects in the source domain [Rosenbaum and Rubin, 1983]. However, Assumption 1 is not enough to identify the causal effects in the target domain. Thus, we further invoke Assumption 2.

Assumption 2 (Transportability). Suppose that

(ii) 0<s(X) ≜ P(G = 1|X)<1 for all X in the source domain, where s(X) is the sampling score.

Assumption 2 is widely adopted in causal effects estimation via data combination in causal inference [Stuart et al., 2011; Hartman et al., 2015; Kern et al., 2016a; Lesko et al., 2017; Buchanan et al., 2018; Li et al., 2022a; Li et al., 2023a; Colnet et al., 2024; Wu et al., 2025; Wu and Mao, 2025]. Assumption 2(ii) indicates that all individuals in the source domain have a positive probability of belonging to the target domain. Assumption 2(i) implies that E[Y (a)|X,G = 1] = E[Y (a)|X,G = 0] = E[Y (a)|X] for a = 0, 1, which ensures the transportability of the CATE from the source domain to the target domain and leads to the identifiability of τ(X), that is,

where the third equality follows from Assumption 1. Thus, under Assumptions 1-2, the reward R(π) can be identified as

Assumption 2 allows the presence of covariate shift, i.e., the distribution of X in the source domain may significantly differ from that in the target domain [Gama et al., 2014].

4 Policy Adaptation Under Covariate Shift

In this section, we proposed a method for learning the optimal policy. It consists of two steps: (a) policy evaluation, estimating the reward R(π) for a given π, and (b) policy learning, learning the optimal policy based on the estimated reward.

4.1 Estimation of the Reward R(π)

According to Eq. (3), a direct method for estimating the reward R(π) is given as

where ˆµa(X) (a = 0, 1) represents the estimated outcome regression function µa(X). This can be implemented by regressing Y on X using the source dataset with A = a. The unbiasedness of the direct estimator ˆRDirect(π) depends on the accuracy of ˆµa(X). If ˆµa(X) is a biased estimator of µa(X), then ˆRDirect(π) will also be a biased estimator of R(π). Moreover, the generalization performance of the direct method is often poor because ˆµa(X) is trained using the source dataset of A = a, but is applied to the whole target dataset [Li et al., 2022b; Wu et al., 2024b]. When there is a significant difference in the covariate distributions between the source and target datasets, the direct method suffers from the problem of model extrapolation, resulting in poor practical performance.

In addition to the direct method, one can use the propensity score e1(X) = P(A = 1|X,G = 1) and sampling score s(X) = P(G = 1|X) to construct the inverse probability weighting (IPW) estimator of R(π). Note that

with ω(X) = (1 - s(X))/s(X). Based on it, the IPW estimator of R(π) is given as

where ˆ e1(X) and ˆ s(X) are estimates of e1(X) and s(X), respectively. The IPW estimator ˆRIPW(π) is an unbiased estimator of R(π) when ˆ e1(X) and ˆ s(X) are accurate estimators of e1(X) and s(X), respectively, i.e., ˆ e1(X) = e1(X) and ˆ s(X) = s(X). However, a limitation of the IPW estimator is its inefficiency, meaning it tends to have a large variance.

The limitations of direct and IPWmethods are essentially caused by the insufficiency of utilizing the observed data. The direct method does not leverage the information of data indicator G and treatment A, while the IPW method does not extract the relationship between covariates X and outcome Y . To fully utilize the observed data, we employ the semiparametric efficiency theory [Tsiatis, 2006] to derive the efficient influence function and the efficiency bound of R(π). This allows us to obtain the semiparametric efficient estimator of R(π). A semiparametric efficient estimator is considered optimal as it reaches the semiparametric efficiency bound, meaning it has the smallest asymptotic variance under several regularity conditions [Newey, 1990].

Theorem 1 (Efficiency Bound of R(π)). Under Assumptions 1-2, the efficient influence function of R(π) is

The semiparametric efficiency bound of R(π) is Var(φR).

Theorem 1 (See Appendix A.1 for proofs) presents the efficient influence function and semiparametric efficiency bound of R(π) under Assumptions 1-2. From Theorem 1, we can construct the semiparametric efficient (SE) estimator of R(π), which is given as

Next, we analyze the theoretical properties of ˆRSE(π).

Proposition 1 (Double Robustness of ˆRSE(π)). Under Assumptions 1 and 2, ˆRSE(π) is an unbiased estimator of R(π) if one of the following conditions is satisfied:

• (i) ˆµa(x) = µa(x), i.e., ˆµa(x) estimates µa(x) accurately for a = 0, 1.

Proposition 1 (See Appendix A.2 for proofs) shows the double robustness of ˆRSE(π), i.e., if the outcome regression function µa(X) for a = 0, 1 can be estimated accurately, or the propensity score e1(X) and the sampling score the s(X) can be estimated accurately, ˆRSE(π) is unbiased estimator of R(π). Compared to the direct method that requires ˆµa(X) = µa(X) for unbiasedness, and the IPWmethod that requires ˆ e1(X) = e1(X) and ˆ s(X) = s(X) for unbiasedness, double robustness provides more reliable results by mitigating the inductive bias caused by inaccurate models for the nuisance parameters e1(X), s(X), and µa(X) for a = 0, 1.

Theorem 2 (Efficiency of ˆRSE(π)). Under the Assumptions 1-2, if ||ˆ e1(x) - e1(x)||2 · ||ˆµa(x) - µa(x)||2 = oP(n-1/2 ) and ||ˆ s(x)-s(x)||2 · ||ˆµa(x)-µa(x)||2 = oP(n-1/2 ) for all x ∈ X and a ∈ {0, 1}, then ˆRSE(π) is a consistent estimator of R(π), and satisfies

where σ2 = V ar(φR) is the semiparametric efficiency bound of R(π), and d-→means convergence in distribution.

Theorem 2 (See Appendix A.3 for proofs) establishes the consistency and asymptotic normality of the proposed estimator ˆRSE(π). In addition, it shows that ˆRSE(π) is semiparametric efficient, i.e., it achieves the semiparametric efficiency bound. These desired properties hold under the mild condition that the nuisance parameters {e(x), s(x), µ0(x), µ1(x)} are estimated at a rate faster than n-1/4 . These conditions are common in causal inference and can be easily satisfied using a variety of flexible machine learning methods [Chernozhukov et al., 2018; Wu et al., 2024a].

4.2 Learning the Optimal Policy

After estimating the reward, we now focus on learning the optimal policy. Recall that for a given hypothesis space Π,

Algorithm 1 Proposed Policy Learning Approach

Input: The source dataset D1 and the target dataset D0. Output: The learned policy ˆπ.

• 1: Stage 1: Fit models ˆµ1(X), ˆµ0(X), ˆ e1(X), ˆ s(X).
• 2: Stage 2:
• 3: while Stop condition is not reached do
• 4: Sample a batch of data fromD0 ∪D1.
• 5: Minimize the loss - ˆR(π, ˆµ0, ˆµ0, ˆ e1, ˆ s,X,A, Y,G) to update π, using the batch sample.
• 6: end while
• 7: Return a learned policy.

the target policy is given as π∗ (x) = argmaxπ∈ΠR(π). Through optimizing different estimators of reward R(π), we obtain different estimator of π∗ , denoted as ˆπ, defined by

where ˆR(π) can be ˆRDirect(π), ˆRIPW(π), or ˆRSE(π). Algorithm 1 summarizes the procedures for learning π∗ .

As discussed in [Athey and Wager, 2021], when learning the policy by optimizing the estimated reward, it achieves better generalization performance if the estimated reward is more efficient. Since the estimator ˆRSE(π) is the most efficient one under Assumptions 1-2, as shown in Theorem 1, we then focus on exploring its finite sample properties and the learned policy obtained by optimizing it. Similar results can also be derived for the direct and IPWmethods. In finite samples, we allow ˆµa(X), ˆ e1(X), and ˆ s(X) to be inaccurate, i.e., they may differ from µa(X), e1(X), and s(X).

The following Proposition 2 presents the bias of ˆRSE(π). Proposition 2 (Bias). Given the learned ˆµa(X) for a = 0, 1, ˆ e1(X), and ˆ s(X), then for any given π, the bias of ˆRSE(π) is

From Proposition 2 (See Appendix A.4 for proofs), the bias of ˆRSE(π) is the product of the estimation errors µa(X) - ˆµa(X)) and s(X)ea(X)(1-ˆ s(X)) - ˆ s(X)ˆ ea(X)(1-s(X)) for a = 0, 1. Clearly, when either ˆµa(X) is close to µa(X), or ˆ s(X) and ˆ e1(X) are close to s(X) and e1(X), ˆRSE(π) will be close to R(π). This further demonstrates the double robustness of ˆR(π).

Next, we show the generalization error bound (or the regret) of the learned policy. For clarity, we define

which is the learned policy by optimizing ˆRSE(π).

Theorem 3 (Generalization Error Bound). For any finite hypothesis space Π, we have that (i) with at least probability 1 - η,

where B(D0, D1, η, Π) equals to

with ˆµAi (Xi) = Ai ˆµ1(Xi) + (1 - Ai)ˆµ0(Xi). (ii) with at least probability 1 - η,

Theorem 3(i) provides the generalization error bound of the learned policy ˆπse, and Theorem 3(ii) presents the difference between the generalization risks of the learned policy and the optimal policy π∗ . Note that when (Yi-ˆµA(Xi))2 are bounded, then B(D0, D1, η, Π) will converge to 0 as the sample size n goes to infinity. Thus, for a sufficiently large sample size n, if the nuisance parameters are estimated with adequate accuracy, the generalization bound of the learned policy will be approximated well by the estimated reward. Additionally, the generalization bound of the learned policy will be close to that of the optimal policy ˆπ∗ .

4.3 Generalizing Policy to the Entire Domain

The approach proposed in Sections 4.1-4.2 is designed to learn the optimal policy in the target domain. In this subsection, we extend the approach to the entire domain, aiming to learn the optimal policy that maximizes V (π). Under Assumption 1-2, the reward V (π) is identified as

Similar to Theorem 1, we present the the efficient influence function and the efficient bound of V (π).

Theorem 4 (Efficiency Bound of V (π)). Under Assumptions 1-2, the efficient influence function of V (π) is

The semiparametric efficiency bound of V (π) is Var(φV ).

From Theorem 4, we can construct the semiparametric efficient (SE) estimator of V (π), which is given as

Similar to Proposition 1 and Theorem 2, it can be shown that ˆ VSE(π) possesses desirable properties, including double robustness, consistency, asymptotic normality, and semiparametric efficiency under regular conditions. Moreover, following the method described in Section 4.2, we can develop an approach to learn the optimal policy in the entire domain based on the estimated reward ˆ VSE(π). For brevity, the detailed description is omitted here.

5 Experiments

We conduct experiments on both simulated datasets and realworld datasets to answer the following questions:

• RQ1: Does the proposed SE method provide a more accurate estimation of the reward?
• RQ2: Does the proposed method learn better policies?
• RQ3: How does the degree of covariate shift affect the performance of the proposed method?

5.1 Experiments on Simulated Datasets

Datasets. We generate the source and target datasets in the simulated experiment. The source dataset consists of 512 individuals. For each individual, the covariates X = (X1,X2,X3)⊤ ∼ N(µ1,Σ1), where µ1 = [10, 3, 7]⊤ is the mean vector and Σ1 ∈ R3×3 is the covariance matrix with the (i, j)-th element being 2-|i-j| for i, j = 1, 2, 3. The treatment A is generated from Bern(0.5), and the potential outcomes are generated as follows: Y (1) = 15 + 0.4 ˜X1 ˜X2 + 0.7 ˜X3 + ϵ, Y (0) = 10 + 0.1 ˜X1 + 0.5 ˜X2 ˜X3 + ϵ, where ˜Xj = Xj · |X|0.1 j + Xj · |X|0.3 j + Xj · |X|0.5 j for j = 1, 2, 3 and ϵ ∼ N(0, 1). The observed outcome is Y = AY (1) + (1 - A)Y (0). The variables X,A, Y are available in the source dataset. For the target dataset, we generate 2,048 individuals. The covariates X ∼ N(µ2,Σ2), where the mean vector µ2 = [9, 4, 6]⊤ and the covariance matrix Σ2 ∈ R3×3 has (i, j)-th element given by 2-|i-j|+1 for i, j = 1, 2, 3. Only X is available in the target dataset. In addition, to assess the performance of the learned policy, we also generate {Y (1), Y (0)} for individuals in the target dataset, using the same method as in the source dataset.

ComparedMethods. The direct and IPWmethod. ˆπdirect = argmaxπ∈Π ˆRDirect(π) and ˆπipw = argmaxπ∈Π ˆRIPW(π).

EvaluationMetrics. The same as [Kitagawa and Tetenov, 2018; Li et al., 2023b], we adopt the metrics below.

• To assess the performance of a policy learning method, we calculate the true reward in the target dataset given by

Also, we define the regret as ∆E = ˆR(π∗ 0 )- ˆR(ˆπ), representing the difference between the reward induced by the oracle policy and that induced by the learned policy ˆπ.

• To evaluate the accuracy of the learned policy ˆπ, we define the policy error as

Table 1: Comparison of estimated rewards, policy errors and welfare changes on the simulated and real-world datasets. The best results are highlighted in bold. RI refers to the relative improvement over the corresponding baseline. SD indicates standard deviation.

the mean square errors between the oracle policy π∗ 0 and ˆπ.

• We define welfare change as

representing the difference between the total rewards induced by ˆπ and the null policy π ≡ 0. A policy learning method is better when it yields a larger true reward, smaller regret, lower policy error, and greater welfare change.

Performance Comparison (RQ1 and RQ2). We average 50 independent trials of policy learning and Table 1 (top panel) reports the average rewards ˆR(ˆπ), regrets ∆E, policy errors, and the welfare changes for different policy learning methods. From it, we have the following observations:

• The proposed SE method achieves the highest reward, smallest regret, lowest policy error, and largest welfare change. Compared to the Direct method (baseline), these improvements are substantial-up to 7.68% in reward, 78.51% in policy error, and 33.43% in welfare changedemonstrating the superiority of the SE method.
• The standard deviations of the Direct and SE methods are significantly smaller than that of the IPWmethod, indicating the instability (large variance) of the IPWmethod.

In summary, the SE method outperforms the competing Direct and IPWmethods, owing to its desirable properties, such as double robustness and semiparametric efficiency.

Effect of varying the degree of covariate shift (RQ3). We evaluate the robustness of three policy learning approaches (Direct, IPW, SE) by varying the degree of covariate shift between the source and target datasets. The degree of covariate shift is measured by the Chebyshev distance between the two mean vectors µ1 and µ2 in source and target datasets. The Chebyshev distance is defined as d(µ1, µ2) = maxj |µ1,j - µ2,j |, where µ1,j and µ2,j are the j-th elements of µ1 and µ2, respectively. Three metrics are used to measure performance: true reward induced by the learned policy, policy error, and welfare change. Figure 1 illustrates how these metrics vary with increasing covariate shift, showing that the proposed SE method remains more stable across different levels of covariate shift and consistently achieves superior performance.

Effect of varying probability of treatment. We further evaluate the robustness of the proposed methods by varying the generation mechanism of treatments in the source dataset. We generate A in the source dataset from Bern(σ(-β ˜X2)), where σ(·) is the sigmoid function, and β = 0, 0.1, . . . , 0.9, 1. The probability of receiving the treatment is 0.5 when β = 0. As β increases, the probability of an individual receiving the treatment decreases (since ˜X2 is always greater than 0 in our setting). Figure 2 displays how the evaluation metrics change as the proportion of individuals receiving treatment decreases, demonstrating the robustness, stability, and overall superiority of the proposed SE method.

5.2 Experiments on Real-World Datasets

Datasets. The Communities and Crime dataset [Redmond, 2009] comprises 1994 records from communities in the United States, which combines socio-economic data from the 1990 US Census, law enforcement data from the 1990 US LEMAS survey, and crime data from the 1995 FBI UCR. Each record includes 127 covariates, including location information (such as state and county), economic factors (such as perCapInc and HousVacant) and demographics (such as PopDens and PctBSorMore). We use records from communities in New Jersey as the source dataset and records from communities in other states as the target dataset. In addition to using the information from covariates, we simulated the treatment A and the potential outcomes Y (1), Y (0). See Appendix D for the detailed data generation process.

Performance Comparison. We also average over 50 independent trials of policy learning and Table 1 (bottom panel) reports the associated results. From it, we have the following observations: (1) The proposed SE method outperforms the Direct and IPWmethods across all evaluation metrics; (2) Compared to the baseline, the SE method shows substantial improvements-up to 26.66% in reward, 84.35% in policy error, and 92.69% in welfare change. These observations further demonstrate the SE method's superiority.

6 RelatedWork

Policy Learning. Policy learning seeks to identify which individuals should receive treatment to maximize the re-

1e5

1e5

ward according to their covariates [Murphy, 2003], with wide-ranging applications in fields such as precision medicine [Bertsimas et al., 2017; Kitagawa and Tetenov, 2018; Kosorok and Laber, 2019], reinforcement learning [Liu et al., 2021; Kwan et al., 2023], and recommender systems [Ma et al., 2020; Chen and Sun, 2021; Li et al., 2023a]. Nevertheless, most policy learning approaches rely solely on a single labeled dataset. In scenarios where labeled data are difficult to obtain, these methods often struggle with external validity and generalizability. To address this issue, a straightforward strategy is to combine labeled and unlabeled data, utilizing techniques like transfer learning and semisupervised learning [Kora et al., 2022; Huynh et al., 2022]. However, while transfer and semi-supervised learning methods are well-developed for prediction tasks, the challenge of learning optimal policies by leveraging both labeled and unlabeled data remains largely unexplored [Uehara et al., 2020].

Causal Effects Generalizability. Recently, there has been growing research interest in integrating information from multiple data sources for causal inference [Hartman et al., 2015; Lodi et al., 2019; Colnet et al., 2024; Wu et al., 2025; Kallus and Mao, 2024; Wu and Mao, 2025]. However, heterogeneity in data distribution across these sources presents a significant challenge. Different heterogeneity in data distribution necessitates tailored techniques, including specific assumptions [Hotz et al., 2005; Kern et al., 2016b; Li et al., 2023c; Hu et al., 2023; Wu et al., 2024c; Yang et al., 2024] and structural causal models (SCM) [Pearl, 1995; Correa et al., 2018; Tikka et al., 2019]. While data integration has gained considerable theoretical attention, its extension to policy learning remains underexplored. In this article, we investigate how to learn the optimal policy in a target dataset by leveraging information from a source dataset.

7 Conclusion

In this article, we propose a principled policy learning method under covariate shift from a causal perspective. We detail the identifiability assumptions for the reward, derive its efficient influence function, and develop a doubly robust estimator. We establish the asymptotic properties of the proposed estimator, including consistency, normality, and semiparametric efficiency. We also provide a generalization error bound for the learned optimal policy. Extensive experiments confirm the effectiveness and reliability of our proposed method, demonstrating both theoretical and practical advantages.

Acknowledgments

This work was supported by the National Key Research and Development Program (No. 2024YFD2100700), the National Natural Science Foundation of China (No. 12301370), the BTBU Digital Business Platform Project by BMEC, the Beijing Key Laboratory of Applied Statistics and Digital Regulation, and the gift funding from ByteDance Research.