DCIL: DropConnect-based Causal Imitation Learning Under Environment Heterogeneity

Abstract

In many real-world scenarios, agents operate across heterogeneous environments where the underlying dynamics and data-generating processes vary. Standard reinforcement learning and imitation learning methods often fail in such settings as they typically assume stationarity and learn policies that overfit to environment-specific correlations. A key challenge is the presence of spurious correlations as observed states often contain both causal and non-causal features, with the latter introducing environment-specific biases that undermine generalization. To address this problem, we propose DropConnect-based Causal Imitation Learning (DCIL), a novel offline imitation learning framework designed to identify and exploit stable causal mechanisms across diverse environments. DCIL introduces a gradient alignment constraint that encourages the policy to align with causal structures shared across training environments. To further mitigate overfitting to spurious correlations, DCIL involves DropConnect-based regularization, injecting stochastic perturbations into network weights to simulate parameter uncertainty and reduce reliance on unstable features. We evaluate DCIL on synthetic benchmarks derived from OpenAI Gym control tasks, where non-causal features exhibiting spurious correlations are explicitly injected to simulate environmental heterogeneity. Experimental results show that DCIL outperforms state-of-the-art imitation learning baselines, achieving superior generalization to unseen environments. These findings highlight the importance of incorporating causal reasoning and structured regularization into policy learning for robust performance under environment shift.