Learning Contextual Causal Dynamics for Robust Exploration in Reinforcement Learning
Abstract
Efficient exploration in sparse-reward environments, where informative extrinsic feedback is scarce, remains a significant challenge in reinforcement learning. A key limitation of existing exploration strategies is that they often assume densely coupled environment dynamics, overlooking the underlying causal mechanisms of the environment, particularly the fact that causal relationships may vary across contexts. To address this issue, we propose contextual causal dynamics model-based intrinsic motivation (CIM), a causality-aware exploration framework that explicitly captures context-dependent sparse causal structures and derives intrinsic motivation signals, specifically causal action influence and curiosity. These intrinsic rewards encourage interventions on causally relevant state factors rather than merely promoting observation novelty-seeking behavior, thereby facilitating robust exploration and active skill acquisition. Moreover, by learning contextual causal mechanisms, CIM identifies and filters out task-irrelevant state variables, improving generalization under distribution shifts. Empirical results demonstrate that our method achieves superior sample efficiency, training robustness, and out-of-distribution generalization performance across multiple robotic manipulation tasks.