mlsanityData to 13 Oct 2026

MoCoDA: Model-based Counterfactual Data Augmentation

arXiv 2210.11287 ↗First posted 20 Oct 2022cs.LG cs.AI cs.ROcited by 44

On the map: Offline RL and imitation learning from data → Offline RL algorithms

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.

Key facts. Cited by 44 papers in the map (data through 13 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.

Builds on

37 of 70 references.

Closest ancestors · 3

Same direction, newest first

  1. 2021A Minimalist Approach to Offline Reinforcement Learning12 Jun
  2. 2020Conservative Q-Learning for Offline Reinforcement Learning8 Jun
  3. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems4 May

Other references · 34

  1. 2014Adam: A Method for Stochastic Optimizationcited by 37,093
  2. 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,656
  3. 2015Continuous control with deep reinforcement learningcited by 3,290
  4. 2018Off-Policy Deep Reinforcement Learning without Explorationcited by 1,231
  5. 2017Hindsight Experience Replaycited by 1,104
  6. 2020Object-Centric Learning with Slot Attentioncited by 869
  7. 2019When to Trust Your Model: Model-Based Policy Optimizationcited by 756
  8. 2018QT-Opt: Scalable Deep Reinforcement Learning for Vision-Based Robotic Manipulationcited by 731
  9. 2018Learning Dexterous In-Hand Manipulationcited by 645
  10. 2020Reinforcement Learning with Augmented Datacited by 456
  11. 2018Neural Relational Inference for Interacting Systemscited by 339
  12. 2019Recurrent Independent Mechanismscited by 219
22 more
  1. 2019Contrastive Learning of Structured World Modelscited by 206
  2. 2018Reducing Gender Bias in Abusive Language Detectioncited by 152
  3. 2019Counterfactual Off-Policy Evaluation with Gumbel-Max Structural Causal Modelscited by 115
  4. 2020Learning to Generalize Across Long-Horizon Tasks from Human Demonstrationscited by 114
  5. 2019Entity Abstraction in Visual Model-Based Reinforcement Learningcited by 110
  6. 2019SCALOR: Generative World Models with Scalable Object Representationscited by 108
  7. 2020Counterfactual Data Augmentation using Locally Factored Dynamicscited by 93
  8. 2018Woulda, Coulda, Shoulda: Counterfactually-Guided Policy Searchcited by 81
  9. 2021Generalization and Robustness Implications in Object-Centric Learningcited by 81
  10. 2022Causal Dynamics Learning for Task-Independent State Abstractioncited by 64
  11. 2020Counterfactual Credit Assignment in Model-Free Reinforcement Learningcited by 61
  12. 2020Sample-Efficient Reinforcement Learning via Counterfactual-Based Data Augmentationcited by 58
  13. 2019Causal Modeling for Fairness in Dynamical Systemscited by 46
  14. 2020A Boolean Task Algebra for Reinforcement Learningcited by 46
  15. 2014Near-optimal Reinforcement Learning in Factored MDPscited by 41
  16. 2020Trust the Model When It Is Confident: Masked Model-based Actor-Criticcited by 36
  17. 2021Fast and Slow Learning of Recurrent Independent Mechanismscited by 36
  18. 2021Action-Sufficient State Representation Learning for Control with Structural Constraintscited by 31
  19. 2019AC-Teach: A Bayesian Actor-Critic Method for Policy Learning with an Ensemble of Suboptimal Teacherscited by 28
  20. 2019Working Memory Graphscited by 23
  21. 2022Policy Architectures for Compositional Generalization in Controlcited by 23
  22. 2020Counterfactually Guided Off-policy Transfer in Clinical Settingscited by 13

Built on by

Cited by year (2026 to 13 Oct)

Citing papers per year: 2023: 10; 2024: 13; 2025: 16; 2026: 5.

44 papers.

Sort:

By area:

  1. 2023MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrationscited by 313
  2. 2025Pseudo-Simulation for Autonomous Drivingcited by 142
  3. 2024DexMimicGen: Automated Data Generation for Bimanual Dexterous Manipulation via Imitation Learningcited by 106
  4. 2024SkillMimicGen: Automated Demonstration Generation for Efficient Skill Learning and Deploymentcited by 49
  5. 2024Dream to Manipulate: Compositional World Models Empowering Robot Imitation Learning with Imaginationcited by 45
  6. 2023What is Essential for Unseen Goal Generalization of Offline Goal-conditioned RL?cited by 36
  7. 2023Seeing is not Believing: Robust Reinforcement Learning against Spurious Correlationcited by 25
  8. 2023Causal Reinforcement Learning: A Surveycited by 22
  9. 2023Interpretable Reward Redistribution in Reinforcement Learning: A Causal Approachcited by 17
  10. 2024RoCoDA: Counterfactual Data Augmentation for Data-Efficient Robot Learning from Demonstrationscited by 13
  11. 2024Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learningcited by 11
  12. 2025Learning Interactive World Model for Object-Centric Reinforcement Learningcited by 11
32 more
  1. 2023Guided Data Augmentation for Offline Reinforcement Learning and Imitation Learningcited by 10
  2. 2024OASIS: Conditional Distribution Shaping for Offline Safe Reinforcement Learningcited by 10
  3. 2024SPARTAN: A Sparse Transformer World Model Attending to What Matterscited by 10
  4. 2024SkiLD: Unsupervised Skill Discovery Guided by Factor Interactionscited by 9
  5. 2023ELDEN: Exploration via Local Dependenciescited by 8
  6. 2023Offline Imitation Learning with Variational Counterfactual Reasoningcited by 7
  7. 2023Understanding when Dynamics-Invariant Data Augmentations Benefit Model-Free Reinforcement Learning Updatescited by 6
  8. 2023MACCA: Offline Multi-agent Reinforcement Learning with Causal Credit Assignmentcited by 6
  9. 2025Null Counterfactual Factor Interactions for Goal-Conditioned Reinforcement Learningcited by 5
  10. 2025Counterfactual Identifiability via Dynamic Optimal Transportcited by 5
  11. 2025Towards Empowerment Gain through Causal Structure Learning in Model-Based RLcited by 4
  12. 2025Curious Causality-Seeking Agents Learn Meta Causal Worldcited by 4
  13. 2024BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learningcited by 3
  14. 2025From Pixels to Factors: Learning Independently Controllable State Variables for Reinforcement Learningcited by 3
  15. 2026A Definition and Roadmap for World Modelscited by 3
  16. 2024Compositional Conservatism: A Transductive Approach in Offline Reinforcement Learningcited by 2
  17. 2024Reinforcement Learning with Euclidean Data Augmentation for State-Based Continuous Controlcited by 2
  18. 2025CausalPlan: Empowering Efficient LLM Multi-Agent Collaboration Through Causality-Driven Planningcited by 2
  19. 2025An Analysis of Causal Effect Estimation using Outcome Invariant Data Augmentationcited by 2
  20. 2024AdaDemo: Data-Efficient Demonstration Expansion for Generalist Robotic Agentcited by 1
  21. 2025GenHPE: Generative Counterfactuals for 3D Human Pose Estimation with Radio Frequency Signalscited by 1
  22. 2025Beyond the Known: Decision Making with Counterfactual Reasoning Decision Transformercited by 1
  23. 2025Individualised Counterfactual Examples Using Conformal Prediction Intervalscited by 1
  24. 2025Learning Causal Structure Distributions for Robust Planningcited by 1
  25. 2025Dynamic Sparsity: Challenging Common Sparsity Assumptions for Learning World Models in Robotic Reinforcement Learning Benchmarkscited by 1
  26. 2024Offline Policy Learning via Skill-step Abstraction for Long-horizon Goal-Conditioned Taskscited by 0
  27. 2025Counterfactual experience augmented off-policy reinforcement learningcited by 0
  28. 2025Q-function Decomposition with Intervention Semantics with Factored Action Spacescited by 0
  29. 2026IPD: Boosting Sequential Policy with Imaginary Planning Distillation in Offline Reinforcement Learningcited by 0
  30. 2026Reflex: Reinforcement Learning with Reflection Symmetry Exploitation in State-Based Continuous Controlcited by 0
  31. 2026MinInter: Minimizing Trajectory Interpolation During Data Augmentation for Imitation Learningcited by 0
  32. 2026Vision-Language Grounded Task-Context-Aware Imitation Learning for Robotic Disassemblycited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning
  2. 2026Low-Rank Dynamics-Effective Latent Carriers for Counterfactual Rollout in Learned World Models
  3. 2023Learning Dynamic Attribute-factored World Models for Efficient Multi-object Reinforcement Learning
  4. 2026Localized Dynamics-Aware Domain Adaption for Off-Dynamics Offline Reinforcement Learning
  5. 2022DARA: Dynamics-Aware Reward Augmentation in Offline Reinforcement Learning
  6. 2025Modeling Unseen Environments with Language-guided Composable Causal Components in Reinforcement Learning
  7. 2021Augmented World Models Facilitate Zero-Shot Dynamics Generalization From a Single Offline Environment
  8. 2026Learning Counterfactual World Models for Embodied Reasoning under Partial Observability