Causal representation learning in offline visual reinforcement learning

Abstract

Real-world reinforcement learning (RL) applications contend with high-dimensional visual observations contaminated by confounding factors, which induce spurious correlations and obscure decision-relevant information. Compounding this issue, the inability to interact online necessitates reliance on pre-collected datasets, thereby hampering a deeper understanding of complex environment structures. In this work, by focusing on the causal rather than spurious correlations in the input and explicitly distinguishing between task-related and task-irrelevant elements of the causal variables, we propose a mask-based algorithm for learning task-related minimal causal state representations, namely MMCS. Specifically, MMCS guides the decoupling of minimal causal variables through mask network partitioning and jointly enforcing conditional independence and causal sufficiency, thereby eliminating unnecessary dependencies between variables and uncovering causal dependency structures. More importantly, MMCS is decoupled from downstream policy learning, and can function as a plug-in method compatible with any offline reinforcement learning algorithm. Empirical results on the Visual-D4RL benchmark demonstrate that MMCS significantly improves performance and sample efficiency in downstream policy learning. In addition, its robust performance in various distraction environments highlights the potential of MMCS to improve the generalizability of offline RL, especially under conditions of limited data and visual distractions. Code is available at https://github.com/DMU-XMU/MMCS.git.