Causal Action Influence Aware Counterfactual Data Augmentation
Núria Armengol Urpí, Marco Bagatella, Marin Vlastelica, Georg Martius
Abstract
Offline data are both valuable and practical resources for teaching robots complex behaviors. Ideally, learning agents should not be constrained by the scarcity of available demonstrations, but rather generalize beyond the training distribution. However, the complexity of real-world scenarios typically requires huge amounts of data to prevent neural network policies from picking up on spurious correlations and learning non-causal relationships. We propose CAIAC, a data augmentation method that can create feasible synthetic transitions from a fixed dataset without having access to online environment interactions. By utilizing principled methods for quantifying causal influence, we are able to perform counterfactual reasoning by swapping action-unaffected parts of the state-space between independent trajectories in the dataset. We empirically show that this leads to a substantial increase in robustness of offline learning algorithms against distributional shift. Videos, code and data are available at https: //sites.google.com/view/caiac .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers9
- Learning Interactive World Model for Object-Centric Reinforcement LearningFan Feng, Phillip Lippe, Sara MagliacaneNeurIPS 2025 · 13 citations
- An Analysis of Causal Effect Estimation using Outcome Invariant Data AugmentationUzair Akbar, Niki Kilbertus, Hao Shen, Krikamol Muandet et al.NeurIPS 2025 · 3 citations
- Dream to Manipulate: Compositional World Models Empowering Robot Imitation Learning with ImaginationLeonardo Barcellona, Andrii Zadaianchuk, Davide Allegro, Samuele Papa et al.ICLR 2025 · 2 citations
- Causal Information Prioritization for Efficient Reinforcement LearningHongye Cao, Fan Feng, Tianpei Yang, Jing Huo et al.ICLR 2025 · 1 citation
- Null Counterfactual Factor Interactions for Goal-Conditioned Reinforcement LearningCaleb Chuck, Fan Feng, Carl Qi, Chang Shi et al.ICLR 2025
Builds on21
- Offline Reinforcement Learning with Implicit Q-LearningIlya Kostrikov, Ashvin Nair, Sergey LevineICLR 2022 · 1,402 citations
- A Minimalist Approach to Offline Reinforcement LearningScott Fujimoto, Shixiang Shane GuNeurIPS 2021 · 1,292 citations
- Object-Centric Learning with Slot AttentionFrancesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran et al.NeurIPS 2020 · 1,275 citations
- Mastering Atari with Discrete World ModelsDanijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, Jimmy BaICLR 2021 · 1,170 citations
- Reinforcement Learning with Augmented DataMichael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto et al.NeurIPS 2020 · 833 citations
Related papers
- Offline Imitation Learning with Variational Counterfactual ReasoningZexu Sun, Bowei He, Jinxin Liu, Xu Chen et al.NeurIPS 2023 · 13 citations
- Diagnosis, Feedback, Adaptation: A Human-in-the-Loop Framework for Test-Time Policy AdaptationAndi Peng, Aviv Netanyahu, Mark K. Ho, Tianmin Shu et al.ICML 2023 · 18 citations
- C2L: Causally Contrastive Learning for Robust Text ClassificationSeungtaek Choi, Myeongho Jeong, Hojae Han, Seung-won HwangAAAI 2022 · 52 citations
- Reward-Preserving Counterfactual State Editing for Offline Reinforcement LearningSiyu Wang, Xiaocong Chen, Mingming Gong, Yong Li et al.ICML 2026
- Causal Flow Q-Learning for Robust Offline Reinforcement LearningMingxuan Li, Junzhe Zhang, Elias BareinboimICML 2026 · 1 citation
