DCIL: DropConnect-based Causal Imitation Learning Under Environment Heterogeneity

1 Introduction

Reinforcement Learning (RL) has achieved impressive success in a wide range of decision-making tasks, including game playing, robotics, and resource management [Arulkumaran et al., 2017; Mnih et al., 2015]. However, many classical RL algorithms [Lillicrap et al., 2015; Hasselt et al.,

2016] assume access to a well-specified and fully observable reward function, which is often impractical to define in real-world applications [Hadfield-Menell et al., 2017]. Imitation Learning (IL) offers a promising alternative by learning policies directly from expert demonstrations, thereby bypassing the need for explicit reward specification. Compared to exploration-based RL, IL is often more sample-efficient and safer, especially in safety-critical domains such as autonomous driving, healthcare, and robot manipulation [Jang et al., 2022; Codevilla et al., 2019; Hu et al., 2022].

Despite these advantages, most IL methods assume that expert demonstrations are drawn from a single stationary environment. In practice, however, expert trajectories often originate from multiple heterogeneous environments, each with different dynamics, feature distributions, or contextual factors. As a result, observed states typically contain both causal features that genuinely influence expert decisions and non-causal features that are spuriously correlated with actions due to environment-specific artifacts [Codevilla et al., 2019]. Policies trained on such data may overfit to environment-dependent correlations, failing to capture the underlying decision-making mechanisms and generalizing poorly to unseen environments. This fundamental mismatch between environment-specific correlations and environmentinvariant decision mechanisms poses a central challenge for offline imitation learning under distribution shift.

Such issues are widespread in real-world settings. For example, in autonomous driving, an agent trained on data collected during rainy weather may observe a strong correlation between windshield wiper usage and reduced speed. However, wiper activity does not causally influence driving policy; it simply correlates with weather conditions, which are the actual causal factor. A policy that mistakes this correlation for causation may make unsafe decisions under different conditions.

These challenges highlight the need for imitation learning algorithms that generalize across environments by identifying and leveraging stable causal mechanisms underlying expert behavior. To this end, we propose DCIL (DropConnectbased Causal Imitation Learning), a novel framework for offline imitation learning from demonstrations collected in heterogeneous environments 1 . DCIL is designed to learn policies that remain robust under distribution shifts without requiring additional online interaction. DCIL incorporates two key components. First, it introduces a gradient alignment constraint that encourages the policy to focus on stable causal features shared across environments. Second, to further reduce sensitivity to spurious features, it applies a DropConnect-based regularization strategy that injects stochastic perturbations into the policy network weights during training. This improves robustness by modeling parameter uncertainty and discouraging reliance on unstable patterns. We validate DCIL on modified OpenAI Gym control tasks where non-causal features are explicitly injected to simulate environmental heterogeneity. Experimental results show that DCIL consistently outperforms state-of-the-art imitation learning methods in generalization to unseen environments. Our contributions are summarized as follows:

1 Code is available at https://github.com/PapriKy/DCIL.

• Address the problem of offline imitation learning from expert trajectories collected across heterogeneous environments, focusing on generalization without additional interaction.
• Propose a novel framework that integrates causal gradient alignment constraints to promote learning of stable, transferable policies.
• Introduce a simple yet effectiveDropConnect-based regularization strategy to alleviate overfitting induced by conventional causal constraints, thereby improving robustness across heterogeneous environments.

2 RelatedWork

Causal inference [Pearl, 2009; Peters et al., 2016] has emerged as a promising direction for improving generalization in decision-making systems, including imitation learning and reinforcement learning. By identifying and leveraging variables that are causally related to decisions, these approaches aim to produce models that remain robust under distributional shifts. De Haan et al. [de Haan et al., 2019] highlight the issue of causal misidentification, where policies inadvertently rely on spurious correlations in expert data rather than true causal variables. This issue becomes more severe when distribution shifts occur during deployment, as noncausal correlations observed during training may no longer hold. To address this problem, their approach uses causal graph-structured policies and targeted interventions, such as expert queries and environment interactions, to infer the correct causal structure. However, such strategies are not applicable in fully offline settings where no additional data collection is permitted. To overcome the need for interventions, Sonar et al. [Sonar et al., 2021] incorporate the principle of Invariant Risk Minimization (IRM) [Arjovsky et al., 2019] into policy learning. This framework encourages policies that perform consistently across multiple environments by aligning invariant components. While effective in certain settings, IRM-based methods can still suffer from spurious correlations even when the regularization penalty approaches zero, potentially leading to overfitting [Lin et al., 2022]. Invariant Causal Imitation Learning (ICIL) [Bica et al., 2021] separates causal and non-causal representations and enforces invariance via adversarial training. However, the combination of mutual information maximization and adversarial objectives often leads to unstable optimization, motivating simpler regularization-based alternatives such as ours. In contrast to these approaches, DCIL avoids explicit adversarial training or causal graph specification, and instead enforces causal invariance through gradient-based constraints combined with stochastic regularization.

3 Methodology

To overcome the limitations of conventional imitation learning, we propose a novel approach that introduces stable causal constraints to disentangle true causal features from environment-specific noise. Our method constrains the gradients of input-layer weights and incorporates DropConnectbased regularizations, jointly capturing stable causal influences across heterogeneous environments while mitigating overfitting to spurious correlations. This mechanism not only reduces reliance on unstable features but also enhances robustness, particularly in settings with limited environment diversity. Unlike prior work, our approach jointly applies a causal constraint that regularizes policy gradients toward stable causal features across environments, and a stochastic, weight-level regularization to mitigate overfitting, resulting in improved generalization in offline imitation learning. An overview of our method is presented in Figure 1.

3.1 Preliminaries

In this section, we formalize the problem of imitation learning across multiple heterogeneous environments using the framework of multi-environment Markov Decision Processes (MDPs). Without loss of generality, we assume that we observe data consisting solely of expert demonstrations collected from a set of K heterogeneous environments. For the environment e ∈ E, where E = {1, . . . , K} and |E| = K, we define the MDP framework as Me = {S e , A, P e , pe0 , re , γ}, where S e denotes the state space specific to the environment e, A represents the action space shared across environments, P e (s′ |s, a) denotes the transition dynamics characterizing the probability of transitioning to state s′ ∈ S e given state s ∈ S e and action a ∈ A, pe0 is the initial state distribution, re (s, a) : S e × A → R is the reward function, and γ ∈ (0, 1) is the discount factor. To further formulate the imitation learning, we also denote ρe as the corresponding occupancy measure of the expert policy, πe , in the environment e. Specifically, we define ρeπ (s) = (1 - γ) ∑∞t=0 γ tP(st = s | πe ) and ρeπ (s, a) = π(a | s)ρeπ (s), which represent the state and state-action visitation distributions under πe . In particular, there exists a bijective correspondence between the policy πe and its occupancy measures ρeπ (s) and ρeπ (s, a).

3.2 Offline Imitation Learning

In offline imitation learning, the agent has no access to environment interaction and the reward function re is unknown. Instead, it is provided with a dataset of expert demonstrations De = {(set , at)}t≥0 collected by an expert policy πe exp in environment e.

We adopt the Maximum Entropy Inverse Reinforcement Learning (MaxEnt IRL) framework [Ho and Ermon, 2016], which formulates imitation learning as a state-action distribution matching problem between the expert and the learned policy. Following IQ-Learn [Garg et al., 2021], the MaxEnt IRL objective can be reformulated into a non-adversarial optimization problem directly in the Q-function space by leveraging the inverse soft Bellman operator.

Specifically, the resulting offline imitation learning objective for environment e is given by

where V Q(s) = log ∑a′ exp(Q(s, a′ )) denotes the soft value function and T π is the inverse soft Bellman operator defined under policy π.

The objective in Eq. (1) encourages the learned policy to match expert behavior using only offline data. However, when expert demonstrations are collected from heterogeneous environments, directly minimizing this objective may lead the policy to exploit spurious correlations between states and actions, resulting in poor generalization. In the following subsection, we introduce causal constraints to mitigate this issue and promote learning of environment-invariant decision mechanisms.

3.3 Stable Causal Feature Learning Under Causal Constraints

In each environment e, the observed state set admits a latent representation that can be decomposed into environmentinvariant causal factors and environment-specific nuisance factors. Specifically, we assume that set is mapped into a latent space Z via a (possibly implicit) representation function, yielding

where ct denotes the latent causal variables that are shared across environments, and µet captures the remaining environment-dependent factors. Motivated by the observation that expert actions often depend only on a subset of taskrelevant features [Peters et al., 2016; Yin et al., 2024], we assume that actions are generated solely from the causal representation, i.e.,

However, since the mapping from the observed state set to its latent decomposition (ct,µet ) is unknown, the underlying causal features governing action generation remain unobserved.

To model these latent causal features and learn a policy that generalizes across environments, we introduce the following assumptions.

Assumption 1 (Heterogeneity). The observed state set can be decomposed into causal variables ct and noncausal variables µet . The non-causal variables are generated via a structural equation (illustrated in Figure 2): µet ← fµ(ct-1,µet-1 , at-1, ξ e ) and Pa(µet ) = {ct-1,µet-1 , at-1, ξ e}, where ξe denotes an environmentspecific exogenous variable that represents a soft or hard intervention on the non-causal variables.

Assumption 2 (Invariance). The causal variables ct are independent of the environment-specific exogenous variable ξ e , i.e., ct ⊥ ξ e . The causal variables are generated via a structural equation (illustrated in Figure 2): ct+1 ← fc(ct, at) with Pa(ct+1) = {ct, at}, and the corresponding transition dynamics across environments satisfy the invariance condition:

Furthermore, the expert policy conditioned on causal states is also invariant across environments:

Assumption 1 facilitates the modeling of spurious correlations by distinguishing invariant causal variables from environment-specific confounders, thereby enabling us to account for heterogeneity across environments. Assumption 2 is crucial for learning a policy that generalizes well. It ensures that both the transition dynamics and expert behavior with respect to the causal variables remain stable across environments, allowing us to apply invariance principles to identify and exploit the shared causal mechanisms. Both assumptions are commonly adopted in invariance-based causal inference [Krueger et al., 2021; Yin et al., 2024].

Based on Assumptions 1 and 2, we consider a representation of the causal variables, i.e., ct = ϕ(set ; θs), to encode the causal features of the action at, where θs ∈ Rm×|S| denotes the input-layer parameters of the policy network with m the number of hidden units and |S| the input dimension. Then, the objective function in Eq. (1) can be adjusted as

Next, we introduce causal constraints for optimization regularization, i.e.,

where ◦ denotes the Hadamard product and ψ(θs) ∈ Rm×|S| denotes the element-wise weighting matrix. Here we use the same weight for all m elements in each column of θs. Specifically, we define the elements in the j-th column of ψ(θs) as ˜wj · ∥θs[:, j]∥2, where θs[:, j] ∈ Rm denotes the j-th column of θs and the scalar ˜wj = 1 ∥θs[:,j]∥2+ε acts as a columnwise normalization factor with ε > 0 a small constant added for numerical stability. Consequently, under Assumptions 1 and 2, the optimality condition of invariant causal features requires the gradients of the imitation learning objective with respect to their associated parameters to vanish across environments. In contrast, parameters corresponding to noncausal features induce environment-dependent gradients and are therefore suppressed through weight shrinkage. By penalizing the weighted gradient norm in Eq. (5), the proposed constraint drives gradients associated with causal features toward zero, while forcing the weights of non-causal features to diminish, thereby enforcing causal invariance at the representation level. By integrating this into Eq. (4), we arrive at the following loss function:

3.4 DropConnect-based Regularization Mechanism

Although causal constraints can encourage invariance across environments, recent studies have shown that such constraints may still lead to spurious solutions and overfitting, even when the corresponding penalty term vanishes [Lin et al., 2022; Zhou et al., 2022]. In practice, enforcing strong causal regularization often introduces additional optimization bias, making the learned policy overly sensitive to specific parameter configurations or training environments. Several existing approaches attempt to mitigate this issue through Bayesian formulations [Lin et al., 2022] or sparsity-inducing regularization [Zhou et al., 2022], which explicitly model uncertainty or restrict model capacity. While effective, these methods typically incur additional modeling complexity or require careful prior specification. In contrast, we adopt a simpler and more practical strategy based on DropConnect. By introducing stochastic perturbations directly at the weight level during training, DropConnect implicitly regularizes the causal constraint, discourages reliance on fragile feature-parameter associations, and effectively alleviates overfitting without altering the underlying optimization objective.

During each forward pass, we introduce a binary mask matrix M ∈ {0, 1}m×|S| , whose elements are independently sampled from a Bernoulli distribution, i.e., Mij ∼ Bernoulli(1 - p), where p is the drop rate hyperparameter. The effective weight matrix used for forward computation is regularized as ˜θs = 1 1-p · (θs ◦ M), which are then used to compute the layer output in the standard manner. Accordingly, the total loss in Eq. (6) can be reformulated as L (˜θs, ˜θπ ):

This stochastic masking mechanism introduces random variations in the model structure during training, encouraging the learned representations to be less dependent on individual connections and more robust to spurious correlations.

4 Experiments

4.1 Experimental Concerns

Our experiments are designed to answer the following questions: Q1: How does DCIL perform in unseen environments? Q2: What is the impact of the number of expert trajectories per environment on the algorithm? Q3: What is the impact of the number of heterogeneous environments on the algorithm? Q4: What is the impact of the dimensionality of non-causal variables on the algorithm? Q5: How effective are the DropConnect-based causal constraints in mitigating the influence of spurious correlations from non-causal variables? Q6: What is the contribution of causal constraints and DropConnect-based regularization to the performance of DCIL? The experiments are organized into 5 parts: OpenAI Gym control tasks, the synthetic GMM dataset [Yin et al., 2024], the visual dataset ColoredMNIST [Arjovsky et al., 2019], the real-world dataset [Beery et al., 2019], and an ablation study.

4.2 Control Tasks FromOpenAI Gym

Benchmark Methods

We compare DCIL with representative offline imitation learning baselines designed for multi-environment or invariant learning settings: (1) BC, which learns policies via supervised imitation; (2) BC-IRM, which combines BC with the IRMv1 objective [Arjovsky et al., 2019] to encourage environment-invariant representations; (3) IQ-Learn [Garg et al., 2021], a non-adversarial distribution-matching approach; (4) IQ-Learn-CoCo, which incorporates the CoCo regularization [Yin et al., 2024] into IQ-Learn to promote invariant representations; and (5) ICIL [Bica et al., 2021], which uses adversarial optimization to identify causal invariance across environments.

Environments and Data Generation

We evaluate DCIL on three OpenAI Gym control tasks: Acrobot, CartPole, and LunarLander. To simulate heterogeneous environments, we augment the original state with additional environment-specific non-causal variables that introduce spurious correlations. Note that the original state may already contain non-causal features, and the augmentation allows us to explicitly control the dimensionality of spurious variables. For each environment e ∈ E, we construct non-causal variables µet ∈ Rdz from the original state features. Specifically, the non-causal variables are generated via an environment-specific linear transformation T e ∈ Rdc×dz , with T e ∼ N(ξ e , 1), where ξ e is an exogenous variable inducing distributional shifts across environments. We compute µet = {c⊤initT e , t = 0, c⊤t-1T e , t > 0, where cinit ∼ U(-1, 1). The aug-

mented state is then given by set = [ct, µet ] ∈ Rdc+dz . When dz = 0, no additional non-causal variables are introduced, corresponding to the original state space, which is evaluated in Figure 5.

(a) Training and Testing Accuracy on GMM

DCIL(ours)

IRM

CoCo REx

ERM

Fishr

(b) Heatmap of Causal Variable Weights on GMM

Figure 6: Performance evaluation on the synthetic GMM task. (a) Training/testing accuracy for CoCo, IRM, ERM, REx, Fishr and our DCIL. (b) Visualization of learned input-layer weights across different methods. The weights related to non-causal features x2 are highlighted by an orange box.

Learning Performance Evaluations

To answer Q1 and Q2, expert demonstrations were collected from two training environments, where the environmentspecific matrices were sampled as T 1 ∼ N(1, 1)dc×dz and T 2 ∼ N(-1, 1)dc×dz , with the non-causal dimensionality set to dz = 6. Test environments were generated by sampling T test ∼ N(4, 1)dc×dz . As shown in Figure 3, the xaxis denotes the number of expert trajectories per training environment, and the y-axis reports the average return of the learned policy, evaluated over 100 randomly sampled test environments. The results show that DCIL achieves near-expert performance using as few as 5 expert trajectories per training environment, and only one trajectory is sufficient in the CartPole task. In contrast, ICIL improves gradually with more demonstrations but fails to reach comparable performance, while BC, BC-IRM, and IQ-Learn remain limited in this setting.

To address Q3, we fixed the number of expert trajectories in each training environment to 5 and set dz = 6. Source environments were constructed by sampling environment-specific matrices from N(1, 1), N(-1, 1), N(0, 1), N(-2, 1), and N(2, 1), while test environments were sampled from N(4, 1)dc×6 . Figure 4 illustrates the effect of the number of training environments on policy performance, where the x-axis indicates the number of source environments and the y-axis shows the average return evaluated over 100 randomly sampled test environments. As the number of training environments increases, the performance of BC and ICIL improves and approaches expert-level performance when trained on five environments. IQ-Learn also benefits from additional environments but remains consistently below the expert policy. BC-IRM exhibits unstable behavior in CartPole. In comparison, DCIL achieves nearexpert performance using only two training environments.

To answer Q4, we collected expert demonstrations from two training environments with 5 expert trajectories each, where T 1 ∼ N(1, 1)dc×dz and T 2 ∼ N(-1, 1)dc×dz . Test environments were sampled from N(4, 1)dc×dz . Figure 5 examines the impact of the non-causal dimensionality dz ∈ {0, 2, 4, 6, 8} on policy learning performance. The xaxis represents different values of dz, and the y-axis shows the average return evaluated over 100 randomly sampled test environments. The results indicate that DCIL consistently maintains near-expert performance as dz increases, whereas the performance of baseline methods becomes increasingly unstable, suggesting that DCIL is robust to spurious correlations introduced by high-dimensional non-causal variables.

Next, to address Q5, we consider a synthetic Gaussian mixture model (GMM) classification task [Yin et al., 2024], ColoredMNIST [Arjovsky et al., 2019], and iWildCam 2019 dataset [Beery et al., 2019]. These benchmarks involve highdimensional and highly non-linear input spaces, and do not rely on imitation learning objectives. This setting allows us to isolate and evaluate the effectiveness of the proposed DropConnect-based causal constraint in suppressing spurious correlations under complex, non-linear representations, independent of the full DCIL framework.

4.3 GaussianMixture Simulation

We study multi-class classification with spurious correlations using a modified Gaussian Mixture Model. For environment e, data are generated via:

where pk = N(xe1 ; µk, I)/ ∑k′ N(xe1 ; µk′ , I) and k1 ∼ Multinomial(1/K, . . . , 1/K). Here xe1 is causal with invariant mapping to ye , while xe2 has spurious environmentdependent correlations with ye . We set K = 5, µk = √1.5Kek ∈ RK , and uek ∼ ∏⌊k/2⌋ i=1 U(0, 1). Training uses pe ∈ {0.01, 0.02}; test uses pf = 0 with new {ufk } to break spurious correlations. For performance comparison, we included several benchmark methods, including CoCo [Yin et al., 2024], IRM [Arjovsky et al., 2019], ERM, REx [Krueger et al., 2021], and Fishr [Rame et al., 2022]. The training/testing classification accuracy and the learned input-layer weights across different methods are illustrated in Figure 6. In particular, the learned weights related to the non-causal variables x2, which exhibit spurious correlations, are highlighted by an orange box in Figure 6 (b). Compared to other benchmark methods, the DropConnect-based causal constraint in DCIL successfully reduces spurious reliance by suppressing weights on x2, promoting better generalization through causal stability.

4.4 ColoredMNIST

ColoredMNIST (CMNIST) [Arjovsky et al., 2019] is a binary classification benchmark where digit colors are spuriously correlated with labels. To evaluate robustness under unknown environments, all training samples are pooled into a single mixed environment without source environment labels. The training data are generated with pe ∈ {0.1, 0.15, . . . , 0.5} and the test environment uses pe = 0.9. Additionally, 10% label noise is introduced. As shown in Table 1, DCIL consistently achieves the best performance across different pe settings, particularly under strong spurious correlations (small pe), demonstrating superior robustness to environmental and label distribution shifts.

4.5 iWildCam 2019 Dataset

To further assess real-world applicability, we additionally conducted experiments on the iWildCam 2019 dataset [Beery et al., 2019], which exhibits substantial distribution shifts across camera traps (e.g., background, illumination, and location). As shown in Table 2, DCIL achieves the best test accuracy (87.28%), outperforming all baseline methods. These results demonstrate that DCIL improves robustness under realistic distribution shifts.

4.6 Ablation Study

Finally, to answer Q6, we conducted ablation studies comparing DCIL with two ablated variants, DCIL[noDrop] and DCIL[noCC]. Specifically, DCIL[noDrop] disables the DropConnect-based regularization while DCIL[noCC] removes the causal constraint component used for identifying stable causal features. In addition, we also included the standard baseline ERM for comparison. Figure 7 visualizes the heatmaps of the first-layer weight matrices θs learned by different methods on the GMM task under varying numbers of training environments (|E| = 2, 3, 4). Each 10×8 matrix represents weights connecting the 8-dimensional input to a hidden layer with 10 units. The first 5 input elements correspond to causal variables, while the remaining 3 elements represent non-causal variables. It can be found that our DCIL outperforms other competing methods and optimizes the weights associated with non-causal inputs (highlighted by orange rectangles) to be closer to 0.

5 Conclusion and Discussion

In this paper, we proposed a novel offline imitation learning framework designed to learn robust and generalizable policies from expert demonstrations collected across heterogeneous environments. To address the challenge of spurious correlations in multi-environment data, we introduced a causal constraint that regularized policy gradients toward stable causal features shared across environments. Additionally, we applied DropConnect-based regularization to inject stochastic perturbations into network parameters during training, which helped reduce overfitting to environment-specific non-causal patterns. This work highlighted the importance of integrating causal reasoning and principled regularization to enhance the robustness and transferability of imitation learning in offline and real-world settings.

Acknowledgments

We thank the anonymous reviewers for their helpful comments on an earlier draft of the paper. Dr. Shihua Li and Xinyu Zhang were supported in this work by the Natural Science Foundation of Jiangsu Province of China [grant number 7708009046A]. Corresponding authors: Dr. Shihua Li and Dr. Rongjie Liu.