PreprintCited by 9Domain Adaptation and Few-Shot LearningMachine Learning and Data ClassificationExplainable Artificial Intelligence (XAI)

MoCoDA: Model-based Counterfactual Data Augmentation

Open access
Pitis, Silviu, Creager, Elliot, Mandlekar, Ajay, Garg, Animesh·2022·arXiv

Abstract

The number of states in a dynamic process is exponential in the number of objects, making reinforcement learning (RL) difficult in complex, multi-object domains. For agents to scale to the real world, they will need to react to and reason about unseen combinations of objects. We argue that the ability to recognize and use local factorization in transition dynamics is a key element in unlocking the power of multi-object reasoning. To this end, we show that (1) known local structure in the environment transitions is sufficient for an exponential reduction in the sample complexity of training a dynamics model, and (2) a locally factored dynamics model provably generalizes out-of-distribution to unseen states and actions. Knowing the local structure also allows us to predict which unseen states and actions this dynamics model will generalize to. We propose to leverage these observations in a novel Model-based Counterfactual Data Augmentation (MoCoDA) framework. MoCoDA applies a learned locally factored dynamics model to an augmented distribution of states and actions to generate counterfactual transitions for RL. MoCoDA works with a broader set of local structures than prior work and allows for direct control over the augmented training distribution. We show that MoCoDA enables RL agents to learn policies that generalize to unseen states and actions. We use MoCoDA to train an offline RL agent to solve an out-of-distribution robotics manipulation task on which standard offline RL algorithms fail.

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2210.11287

Cites

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • It can be difficult to generate plausible interactions without prior knowledge of physics or causal dependencies, especially for new scenes, objects, or robots.

    “however, it can be difficult to generate plausible interactions without prior knowledge of physics [35] or causal dependencies [41,42], especially for new scenes, objects, or robots.”

    From MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations (2023) · cites this paper as [42]

    Section 4 (p. 6) says generating augmented transitions requires known causal parent sets, and Section 2.2 (p. 4) calls discovering dynamics factorization an open problem. This supports the difficulty of generating transitions without causal-dependency information, but the paper does not specifically establish that plausible interactions are especially difficult for new scenes, objects, or robots. (LLM, conf=0.85)

    Experimental — accuracy not yet measured
  • Because data augmentation is offline, it can be difficult to generate plausible interactions without prior knowledge of physics or causal dependencies, especially for new scenes, objects, or robots.

    “However, because data augmentation is offline, it can be difficult to generate plausible interactions without prior knowledge of physics [35] or causal dependencies [41,42], especially for new scenes, objects, or robots.”

    From MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations (2023) · cites this paper as [42]

    Section 4 (p. 6) specifies known parent sets for generating augmented transitions, and Section 2.2 (p. 4) calls learning the dynamics factorization an open problem. This supports the difficulty of augmentation without prior causal knowledge, but the paper does not specifically establish that plausible interactions are especially difficult to generate for new scenes, objects, or robots. (LLM, conf=0.82)

    Experimental — accuracy not yet measured
  • Offline data augmentation can be effective for generating larger dataset for robot manipulation.

    “Offline data augmentation can be effective for generating larger dataset for robot manipulation [7, 35-45]”
  • Several works have used offline data augmentation to increase the dataset size for learning policies.

    “Several works have used offline data augmentation to increase the dataset size for learning policies [7, 35-45].”

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

Still being checked.

Research works printed but not linked: 56

  1. [2][2] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning , pages 104-114. PMLR, 2020. -
  2. [3][3] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. In Advances in neural information processing systems , pages 5048-5058, 2017. -
  3. [4][4] OpenAI: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research , 39(1):3-20, 2020. -
  4. [5][5] Bharathan Balaji, Petros Christodoulou, Xiaoyu Lu, Byungsoo Jeon, and Jordan Bell-Masterson. Factoredrl: Leveraging factored graphs for deep reinforcement learning. In NeurIPS workshop on Offline RL , 2020. -
  5. [7][7] Lars Buesing, Theophane Weber, Yori Zwols, Sebastien Racaniere, Arthur Guez, Jean-Baptiste Lespiau, and Nicolas Heess. Woulda, coulda, shoulda: Counterfactually-guided policy search. International Conference on Learning Representations , 2019. -
  6. [8][8] Guillaume Chaslot, Sander Bakkes, Istvan Szita, and Pieter Spronck. Monte-carlo tree search: A new framework for game ai. In Proceedings of the AAAI Conference on Artificial Intelligence and Interactive Digital Entertainment , volume 4, pages 216-217, 2008. -
  7. [9][9] Elliot Creager, David Madras, Toniann Pitassi, and Richard Zemel. Causal modeling for fairness in dynamical systems. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning , volume 119 of Proceedings of Machine Learning Research , pages 2185-2195. PMLR, 13-18 Jul 2020. -
  8. [10][10] Alexander D'Amour, Hansa Srinivasan, James Atwood, Pallavi Baljekar, David Sculley, and Yoni Halpern. Fairness is not static: deeper understanding of long term fairness via simulation studies. In Proceedings of the 2020 Conference on Fairness, Accountability, and Transparency , pages 525-534, 2020. -
  9. [11][11] Pieter-Tjerk De Boer, Dirk P Kroese, Shie Mannor, and Reuven Y Rubinstein. A tutorial on the crossentropy method. Annals of operations research , 134(1):19-67, 2005. -
  10. [12][12] Andrea Dittadi, Samuele Papa, Michele De Vita, Bernhard Schölkopf, Ole Winther, and Francesco Locatello. Generalization and robustness implications in object-centric learning. In Proceedings of the international conference on Machine learning , 2022. -
  11. [13][13] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems , 34, 2021. -
  12. [14][14] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning , pages 2052-2062. PMLR, 2019. -
  13. [15][15] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations , 2021. -
  14. [16][16] Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research , 19:399-468, 2003. -
  15. [17][17] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870. PMLR, 2018. -
  16. [18][18] Lily Hu and Issa Kohler-Hausmann. What's sex got to do with fair machine learning? In Proceedings of the 2020 Conference on Fairness, Accountability, and Transparency , 2020. -
  17. [19][19] Biwei Huang, Chaochao Lu, Liu Leqi, José Miguel Hernández-Lobato, Clark Glymour, Bernhard Schölkopf, and Kun Zhang. Action-sufficient state representation learning for control with structural constraints. In International Conference on Machine Learning , pages 9260-9279. PMLR, 2022. -
  18. [20][20] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems , pages 12498-12509, 2019. -
  19. [21][21] Edwin T Jaynes. Information theory and statistical mechanics. Physical review , 106(4):620, 1957. -
  20. [22][22] Jindong Jiang, Sepehr Janghorbani, Gerard De Melo, and Sungjin Ahn. Scalor: Generative world models with scalable object representations. In International Conference on Learning Representations , 2019. -
  21. [23][23] Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning , pages 651-673. PMLR, 2018. -
  22. [24][24] Michael Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In IJCAI , volume 16, 6. pages 740-747, 1999. -
  23. [25][25] Taylor W Killian, Marzyeh Ghassemi, and Shalmali Joshi. Counterfactually guided policy transfer in clinical settings. In Conference on Health, Inference, and Learning , pages 5-31. PMLR, 2022. -
  24. [26][26] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014. - · arXiv
  25. [27][27] Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning , pages 2688-2697. PMLR, 2018. -
  26. [28][28] Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In International Conference on Learning Representations , 2020. -
  27. [29][29] Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques . MIT press, 2009. -
  28. [30][30] Emmanouil Krasanakis, Eleftherios Spyromitros-Xioufis, Symeon Papadopoulos, and Yiannis Kompatsiaris. Adaptive sensitive reweighting to mitigate bias in fairness-aware classification. In Proceedings of the 2018 world wide web conference , pages 853-862, 2018. -
  29. [31][31] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:1179-1191, 2020. -
  30. [32][32] Andrey Kurenkov, Ajay Mandlekar, Roberto Martin-Martin, Silvio Savarese, and Animesh Garg. Ac-teach: A bayesian actor-critic method for policy learning with an ensemble of suboptimal teachers. In Conference on Robot Learning , pages 717-734. PMLR, 2020. -
  31. [33][33] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In Advances in Neural Information Processing Systems , 2020. -
  32. [34][34] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020. - · arXiv
  33. [36][36] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems , 33:11525-11538, 2020. -
  34. [37][37] Ricky Loynd, Roland Fernandez, Asli Celikyilmaz, Adith Swaminathan, and Matthew Hausknecht. Working memory graphs. In International conference on machine learning , pages 6404-6414. PMLR, 2020. -
  35. [38][38] Chaochao Lu, Biwei Huang, Ke Wang, José Miguel Hernández-Lobato, Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual-based data augmentation. In NeurIPS Workshop on Offline Reinforcement Learning , 2020. -
  36. [39][39] Kanika Madan, Nan Rosemary Ke, Anirudh Goyal, Bernhard Schölkopf, and Yoshua Bengio. Fast and slow learning of recurrent independent mechanisms. In International Conference on Learning Representations , 2021. -
  37. [41][41] Thomas Mesnard, Theophane Weber, Fabio Viola, Shantanu Thakoor, Alaa Saade, Anna Harutyunyan, Will Dabney, Thomas S Stepleton, Nicolas Heess, Arthur Guez, et al. Counterfactual credit assignment in model-free reinforcement learning. In International Conference on Machine Learning , pages 7654-7664. PMLR, 2021. -
  38. [42][42] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015. -
  39. [43][43] Geraud Nangue Tasse, Steven James, and Benjamin Rosman. A boolean task algebra for reinforcement learning. Advances in Neural Information Processing Systems , 33:9497-9507, 2020. -
  40. [44][44] Michael Oberst and David Sontag. Counterfactual off-policy evaluation with gumbel-max structural causal models. In International Conference on Machine Learning , pages 4881-4890. PMLR, 2019. -
  41. [45][45] Ian Osband and Benjamin Van Roy. Near-optimal reinforcement learning in factored MDPs. Advances in Neural Information Processing Systems , 27, 2014. -
  42. [46][46] Feiyang Pan, Jia He, Dandan Tu, and Qing He. Trust the model when it is confident: Masked model-based actor-critic. Advances in neural information processing systems , 33:10537-10546, 2020. -
  43. [47][47] Ji Ho Park, Jamin Shin, and Pascale Fung. Reducing gender bias in abusive language detection. In Conference on Empirical Methods in Natural Language Processing , 2018. -
  44. [48][48] Judea Pearl. Probabilities of causation: Three counterfactual interpretations and their identification. Synthese , pages 93-149, 1999. -
  45. [49][49] Judea Pearl. Causality . Cambridge university press, 2009. -
  46. [51][51] Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems , 33:3976-3990, 2020. -
  47. [52][52] Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons, 2014. -
  48. [53][53] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature , 550(7676):354-359, 2017. -
  49. [54][54] Shagun Sodhani, Sergey Levine, and Amy Zhang. Improving generalization with approximate factored value functions. In ICLR Workshop on the Elements of Reasoning: Objects, Structure and Causality , 2022. URL https://openreview.net/forum?id=B4exBrOUceq . -
  50. [55][55] Alexander L Strehl. Model-based reinforcement learning in factored-state mdps. In 2007 IEEE International Symposium on Approximate Dynamic Programming and Reinforcement Learning , pages 103-110. IEEE, 2007. -
  51. [56][56] Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin , 2(4):160-163, 1991. -
  52. [57][57] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018. -
  53. [58][58] Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. In Conference on Robot Learning , pages 1439-1456. PMLR, 2020. -
  54. [59][59] Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International conference on learning representations , 2018. -
  55. [60][60] Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for taskindependent state abstraction. In International Conference on Machine Learning , pages 23151-23180. PMLR, 2022. -
  56. [62][62] Allan Zhou, Vikash Kumar, Chelsea Finn, and Aravind Rajeswaran. Policy architectures for compositional generalization in control. arXiv preprint arXiv:2203.05960 , 2022. · arXiv

Cited by

Sorted by citations

Authors
Pitis, S., Creager, E., Mandlekar, A., & Garg, A.
Venue
arXiv
Publisher
arXiv
Published
2022
Pages
18143-18156
ISSN
1049-5258 (ISSN-L)
OpenAlex
W7133200833
Cited by
8
References
0
Citations by year
2026: 4 · 2025: 2 · 2024: 2
FWCI
0.55
Type
conference paper
Language
en
Open access
Locations
  • neural information processing systems (publishedVersion)
  • arXiv (Cornell University) (submittedVersion)
  • arXiv (Cornell University)
Funders
  • Canadian Institute for Advanced Research
  • Government of Canada
  • University of Toronto
  • Natural Sciences and Engineering Research Council of Canada
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Domain Adaptation and Few-Shot Learning
  • Physical Sciences › Computer Science › Artificial Intelligence › Machine Learning and Data Classification
  • Physical Sciences › Computer Science › Artificial Intelligence › Explainable Artificial Intelligence (XAI)
Keywords
counterfactual data augmentation, counterfactual reasoning, data augmentation

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review