# MoCoDA: Model-based Counterfactual Data Augmentation

arXiv: https://arxiv.org/abs/2210.11287 · HTML: https://mlsanity.com/paper/2210.11287/

## Key facts

- Cited by 44 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 37 of its 70 references that are in the map.
- First posted on arXiv 20 Oct 2022.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 30 of 37)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulation](https://mlsanity.com/paper/1806.10293/) (2018)
- [Learning Dexterous In-Hand Manipulation](https://mlsanity.com/paper/1808.00177/) (2018)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Neural Relational Inference for Interacting Systems](https://mlsanity.com/paper/1802.04687/) (2018)
- [Recurrent Independent Mechanisms](https://mlsanity.com/paper/1909.10893/) (2019)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Reducing Gender Bias in Abusive Language Detection](https://mlsanity.com/paper/1808.07231/) (2018)
- [Counterfactual Off-Policy Evaluation with Gumbel-Max Structural Causal Models](https://mlsanity.com/paper/1905.05824/) (2019)
- [Learning to Generalize Across Long-Horizon Tasks from Human Demonstrations](https://mlsanity.com/paper/2003.06085/) (2020)
- [Entity Abstraction in Visual Model-Based Reinforcement Learning](https://mlsanity.com/paper/1910.12827/) (2019)
- [SCALOR: Generative World Models with Scalable Object Representations](https://mlsanity.com/paper/1910.02384/) (2019)
- [Counterfactual Data Augmentation using Locally Factored Dynamics](https://mlsanity.com/paper/2007.02863/) (2020)
- [Woulda, Coulda, Shoulda: Counterfactually-Guided Policy Search](https://mlsanity.com/paper/1811.06272/) (2018)
- [Generalization and Robustness Implications in Object-Centric Learning](https://mlsanity.com/paper/2107.00637/) (2021)
- [Causal Dynamics Learning for Task-Independent State Abstraction](https://mlsanity.com/paper/2206.13452/) (2022)
- [Counterfactual Credit Assignment in Model-Free Reinforcement Learning](https://mlsanity.com/paper/2011.09464/) (2020)
- [Sample-Efficient Reinforcement Learning via Counterfactual-Based Data Augmentation](https://mlsanity.com/paper/2012.09092/) (2020)
- [A Boolean Task Algebra for Reinforcement Learning](https://mlsanity.com/paper/2001.01394/) (2020)
- [Causal Modeling for Fairness in Dynamical Systems](https://mlsanity.com/paper/1909.09141/) (2019)
- [Near-optimal Reinforcement Learning in Factored MDPs](https://mlsanity.com/paper/1403.3741/) (2014)

## Built on by (most cited first, 30 of 44)

- [MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations](https://mlsanity.com/paper/2310.17596/) (2023)
- [Pseudo-Simulation for Autonomous Driving](https://mlsanity.com/paper/2506.04218/) (2025)
- [DexMimicGen: Automated Data Generation for Bimanual Dexterous Manipulation via Imitation Learning](https://mlsanity.com/paper/2410.24185/) (2024)
- [SkillMimicGen: Automated Demonstration Generation for Efficient Skill Learning and Deployment](https://mlsanity.com/paper/2410.18907/) (2024)
- [Dream to Manipulate: Compositional World Models Empowering Robot Imitation Learning with Imagination](https://mlsanity.com/paper/2412.14957/) (2024)
- [What is Essential for Unseen Goal Generalization of Offline Goal-conditioned RL?](https://mlsanity.com/paper/2305.18882/) (2023)
- [Seeing is not Believing: Robust Reinforcement Learning against Spurious Correlation](https://mlsanity.com/paper/2307.07907/) (2023)
- [Causal Reinforcement Learning: A Survey](https://mlsanity.com/paper/2307.01452/) (2023)
- [Interpretable Reward Redistribution in Reinforcement Learning: A Causal Approach](https://mlsanity.com/paper/2305.18427/) (2023)
- [RoCoDA: Counterfactual Data Augmentation for Data-Efficient Robot Learning from Demonstrations](https://mlsanity.com/paper/2411.16959/) (2024)
- [Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning](https://mlsanity.com/paper/2406.03234/) (2024)
- [Learning Interactive World Model for Object-Centric Reinforcement Learning](https://mlsanity.com/paper/2511.02225/) (2025)
- [Guided Data Augmentation for Offline Reinforcement Learning and Imitation Learning](https://mlsanity.com/paper/2310.18247/) (2023)
- [OASIS: Conditional Distribution Shaping for Offline Safe Reinforcement Learning](https://mlsanity.com/paper/2407.14653/) (2024)
- [SPARTAN: A Sparse Transformer World Model Attending to What Matters](https://mlsanity.com/paper/2411.06890/) (2024)
- [SkiLD: Unsupervised Skill Discovery Guided by Factor Interactions](https://mlsanity.com/paper/2410.18416/) (2024)
- [ELDEN: Exploration via Local Dependencies](https://mlsanity.com/paper/2310.08702/) (2023)
- [Offline Imitation Learning with Variational Counterfactual Reasoning](https://mlsanity.com/paper/2310.04706/) (2023)
- [Understanding when Dynamics-Invariant Data Augmentations Benefit Model-Free Reinforcement Learning Updates](https://mlsanity.com/paper/2310.17786/) (2023)
- [MACCA: Offline Multi-agent Reinforcement Learning with Causal Credit Assignment](https://mlsanity.com/paper/2312.03644/) (2023)
- [Null Counterfactual Factor Interactions for Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2505.03172/) (2025)
- [Counterfactual Identifiability via Dynamic Optimal Transport](https://mlsanity.com/paper/2510.08294/) (2025)
- [Towards Empowerment Gain through Causal Structure Learning in Model-Based RL](https://mlsanity.com/paper/2502.10077/) (2025)
- [A Definition and Roadmap for World Models](https://mlsanity.com/paper/2607.06401/) (2026)
- [Curious Causality-Seeking Agents Learn Meta Causal World](https://mlsanity.com/paper/2506.23068/) (2025)
- [BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2407.10967/) (2024)
- [Reinforcement Learning with Euclidean Data Augmentation for State-Based Continuous Control](https://mlsanity.com/paper/2410.12983/) (2024)
- [From Pixels to Factors: Learning Independently Controllable State Variables for Reinforcement Learning](https://mlsanity.com/paper/2510.02484/) (2025)
- [Compositional Conservatism: A Transductive Approach in Offline Reinforcement Learning](https://mlsanity.com/paper/2404.04682/) (2024)
- [CausalPlan: Empowering Efficient LLM Multi-Agent Collaboration Through Causality-Driven Planning](https://mlsanity.com/paper/2508.13721/) (2025)

## Abstract (arXiv)

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.