Causal Dynamics Learning with Quantized Local Independence Discovery

Incorporating causal relationships between the variables into dynamics learning has emerged as a promising approach to enhance robustness and generalization in reinforcement learning (RL). Recent studies have focused on inferring the causal structure of the transition dynamics and leveraging only relevant subsets of the state and action variables for prediction or counterfactual reasoning. However, such approaches tend to overlook the fine-grained local independence relationships that exist among variables. In this work, we propose a novel approach to dynamics learning which infers event-specific causal relationships that hold under certain circumstances referred to as events. Our main idea is to learn a discrete latent variable that represents both the events and corresponding local causal structures via vector quantization. Compared to the prior models using the global causal structure, our approach provides a more detailed understanding of the dynamics by capturing event-specific causal relationships and locally invariant causal mechanisms. Experimental results demonstrate that our method successfully discovers event-specific causal structures, is robust to locally spurious correlations, and generalizes well to downstream tasks compared to previous approaches.