PaperlayerView exact revision on arXiv

Source-linked AI summary

MoCoDA: Model-based Counterfactual Data Augmentation

Silviu Pitis, Elliot Creager, Ajay Mandlekar, Animesh Garg

arXiv:2210.11287v1cs.LGcs.AIcs.RO

TL;DR

Offline and zero-shot RL may lack experience covering states needed for a target task. MoCoDA uses a locally factored dynamics model and an augmented state-action distribution to generate training transitions; the paper reports generalization to unseen states and success on an out-of-distribution robotics task where standard offline RL fails.

  • Problem

    Offline and zero-shot RL may use data unrelated to the target task and lacking coverage of possible states.

    Source: §1 · p. 1

  • Method

    MoCoDA controls the training distribution with a parent distribution and locally factored dynamics to generate augmented transitions for a downstream RL agent.

    Source: §4 · p. 6

  • Results

    MoCoDA-trained offline RL solves an out-of-distribution robotics manipulation task where standard offline RL algorithms fail, and generalizes to states the agent has never seen.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The generated single-step transitions are compatible with any downstream RL algorithm that consumes single-step transitions.

    Source: §6 · p. 10

    Learning locally factored representations, especially from high-dimensional inputs such as pixels, remains future work.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Offline and zero-shot RL can lack coverage of states needed for a target task. MoCoDA uses locally factored dynamics to generate controllable counterfactual training data, and demonstrates generalization to unseen states and an out-of-distribution robotics task.

1 Introduction

  • Offline and zero-shot settings remain challenging when training data are unrelated to the target task and lack coverage of possible states.Source: §1 · p. 1
  • MoCoDA uses a locally factored transition model to generate counterfactual training distributions, including states needed for policies whose optimal behavior requires unseen states.The augmented state-action distribution can be focused on areas relevant to the target task.Source: §1 · pp. 1–2
  • The proposed masked dynamics model supports data augmentation in locally factored settings, relaxing strong assumptions made by prior work on factored MDPs and counterfactual augmentation.Source: §1 · p. 2
  • Controlling the state-action training distribution in a task-relevant way can improve performance.Source: §1 · p. 2
  • MoCoDA-trained offline RL solves an out-of-distribution robotics manipulation task where standard offline RL algorithms fail.The policy also generalizes zero-shot to states it has never seen.Source: §1 · p. 2

2 Preliminaries

The preliminaries distinguish globally factored transition models from local causal models, where each neighborhood has its own factorization. Related work motivates this local view and positions MoCoDA against data augmentation and structured-dynamics approaches.

2.1 Background

  • Factored MDPs assume a global factorization in which each state variable depends on a subset of prior state-action variables.These parent sets define causal mechanisms.Source: §2.1 · p. 3
  • A local causal model partitions the state-action space into neighborhoods, each with its own factored transition function; unlike a factored MDP, the factorization can differ by neighborhood.The framework assumes each graphical structure uniquely identifies its neighborhood.Source: §2.1 · p. 3

2.2 Related Work

  • CoDA stitches empirical samples from disconnected causal mechanisms but requires complete disentanglement; MoCoDA's learned model supports overlapping parent sets and lets agents control the augmented data distribution.Source: §2.1 · p. 3
  • Prior structured-dynamics work often improves model architecture, while MoCoDA also uses structure for data augmentation; learning entity representations and dynamics factorization remain open, orthogonal problems.The cited causal-dynamics approach models variables as unstructured vectors and samples whole mechanisms rather than action vectors alone.Source: §2.2 · p. 4

3 Generalization Properties of Locally Factored Models

Known local factorization can sharply reduce dynamics-model sample complexity and support generalization to unseen state-action pairs. For value functions and policies, the paper instead proposes generating out-of-distribution training data with the locally factored model, because long-horizon predictions can become entangled.

3.1 Sample Complexity of Training a Locally Factored Dynamics Model

  • O(mk^2(|S|^2|A|/m^2)^(1/k) log(|S||A|)) replaces O(|S|^2|A| log(|S||A|)) for full-state model learning under local factorization.The theoretical discussion concerns tabular RL, while the authors apply their method to continuous domains.Source: §3.1 · pp. 4–5
  • For large enough state-action spaces, the exponentially smaller sample count implies the locally factored model must generalize to unseen states and actions.The passage contrasts this sample count with the Ω(|S||A|) count that would be needed without such generalization.Source: §3.1 · p. 5

3.2 Training Value Functions and Policies for Out-of-Distribution Generalization

  • Long-horizon value and policy predictions may lose local sparsity as causal dependencies compound over time, making direct generalization from empirical data difficult.The paper notes that future interactions, such as pool-ball collisions, can make planning depend on the relative positions of all balls, even with factored rewards.Source: §3.2 · p. 5
  • The authors hypothesize rapid temporal entanglement in complex natural systems, while acknowledging that long-horizon structure can remain sparse and allow out-of-distribution generalization.They therefore do not claim that value functions and policies can never generalize well out of distribution.Source: §3.2 · p. 5
  • MOCODA generates out-of-distribution states and actions, applies the locally factored dynamics model to create transitions, and trains value and policy modules on the resulting data.The framework combines the augmented transitions with empirical data, labels both using the target task reward, and feeds them to the RL algorithm.Source: §3.2 · p. 5, Figure 3 · p. 6

4 Model-based Counterfactual Data Augmentation

MOCODA uses a parent distribution and a learned dynamics model to generate augmented transitions for training an RL agent. Its distribution can be controlled to target useful out-of-distribution data, while model and algorithm choices depend on the task setting.

4 Model-based Counterfactual Data Augmentation

  • MOCODA samples a parent distribution, generates transitions with a learned dynamics model, and trains an off-policy RL agent on the augmented data, optionally alongside experience data.The framework controls the training distribution through the parent distribution and dynamics model, both trained using experience data.Source: §4 · p. 6
  • Unlike CoDA, MOCODA permits overlapping parent sets and direct control over the parent distribution.CoDA swaps non-overlapping parent sets and does not allow parent-distribution control.Source: §4 · p. 6

4.1 Generating the Parent Distribution

  • EMP data and short-horizon DYNA rollouts cover only a small neighborhood, while RAND offers no guarantee that the model generalizes to sampled state-actions.The model is expected to generalize when each parent-set configuration has sufficient support in the empirical data.Source: §4.1 · pp. 6–7
  • MOCODA matches each local parent set's augmented marginal to its empirical marginal and selects the maximum-entropy distribution among those matches.The paper describes sampling parent sets sequentially, conditioning on previous partial samples when parent sets overlap.Source: §4.1 · p. 7
  • MOCODA can preserve compounded parent-set imbalances and waste training on irrelevant or overly broad support; MOCODA-U rebalances, while MOCODA-P prioritizes task-relevant regions.MOCODA-U can be impractical with many parent sets; MOCODA-P uses additional task knowledge and has support within MOCODA.Source: §4.1 · p. 7

4.2 The Choice of Dynamics Model and RL Algorithm

  • The experiments use a locally factored model when local structure is known, then label generated transitions with target-task rewards and train an RL algorithm suited to the setting.The authors note that a locally factored model may not be strictly necessary if the underlying dynamics are factored; offline algorithms are used for the offline experiments.Source: §4.2 · pp. 7–8

5 Experiments

Across 2D Navigation and HookSweep2, MOCODA-augmented data improves offline RL outcomes on tasks requiring out-of-distribution generalization. Locally factored dynamics models also show less out-of-distribution error degradation than the compared alternatives, with the stated exception of EMP.

5 Experiments

  • The locally factored model performs better on every tested distribution except EMP and degrades less out-of-distribution than the compared models.The comparison used 2D Navigation dynamics models trained on 35,000 transitions; the passage notes that random-action DYNA data can be out-of-distribution for every model.Source: Table 1 · p. 9, §5 · p. 9
  • In 2D Navigation, augmented MOCODA and MOCODA-U datasets improve average steps across algorithms, supporting learning the diagonal route beyond EMP coverage.The reported results average the last 50 training epochs; RAND also improves performance according to the table caption.Source: Table 2 · p. 9, §5 · p. 9
  • For HookSweep2, EMP misses the goal regions, whereas MOCODA provides some success examples and MOCODA-P reweights block-position combinations for more even table coverage.The prioritized parent distribution uses knowledge that the task involves placing two blocks.Source: §5 · p. 9, Figure 5 · p. 10
  • MOCODA and MOCODA-U improve 2D Navigation results across algorithms, while MOCODA is needed for HookSweep2 learning and MOCODA-P improves its success results.The 2D agents learn the diagonal policy despite its being out-of-distribution relative to EMP.Source: Table 2 · p. 9, §5 · p. 9, Table 3 · p. 10

6 Conclusion

The conclusion presents MOCODA as a framework for augmenting insufficiently covered data with a controllable parent distribution and locally factored dynamics model. The authors report that augmented samples help RL agents learn policies that traverse unseen states and actions, while identifying future research directions and fairness considerations.

6 Conclusion

  • MOCODA adds augmented samples that help RL agents learn policies traversing states and actions absent from the experience buffer.The generated transitions are compatible with any downstream RL algorithm that consumes single-step transitions.Source: §6 · p. 10
  • The conclusion identifies learning locally factorized representations, predicting out-of-distribution generalization, building entity-based RL benchmarks, and re-balancing downstream training data as future directions.It notes that learning such representations in high-dimensional inputs, including pixels, is one direction.Source: §6 · p. 10
  • The authors caution that asserting independence relationships and re-balancing datasets may have fairness implications or result in data biases.They say re-balancing schemes should be sensitive to these possibilities.Source: §6 · p. 10

Checklist

The checklist records the paper’s responses on limitations, societal impacts, theoretical assumptions and proofs, and reproducibility materials.

Checklist

  • The authors report describing the work’s limitations and discussing potential negative societal impacts, marking both checklist items “[Yes]”; the limitations discussion is identified as appearing in Subsection 4.1.Source: Checklist · p. 14
  • The authors report stating the assumptions and providing complete proofs for theoretical results in Appendix A.Source: Checklist · p. 14
  • The authors report providing code, data, and reproduction instructions for the main experimental results.The checklist gives the project repository as the location.Source: Checklist · p. 14

A Proof of Theorem 1

The proof estimates each local causal mechanism and combines the resulting errors to bound the aggregate transition model. It then relates the per-mechanism sample requirement to total state-action visits.

A Proof of Theorem 1

  • Each local causal mechanism has separate count-based parameters for possible child outcomes at each parent configuration.Source: A Proof of Theorem 1 · p. 15
  • The proof uses a concentration bound for empirical estimates of discrete distributions and a lemma that aggregates per-mechanism errors into a transition-model error bound.Source: A Proof of Theorem 1 · p. 15
  • A union bound across states and actions completes the error argument, after which the proof rearranges the bound to obtain the per-mechanism sample requirement.Source: A Proof of Theorem 1 · p. 15
  • The corollary translates per-mechanism samples into total state-action visits, while the authors leave value-difference bounds via the Simulation Lemma or related work for future theory.Source: A Proof of Theorem 1 · pp. 15–16

B Implementation Details

The appendix outlines implementation choices for causal masks, parent-distribution sampling, and dynamics models, directing readers to the code for precise details.

B Implementation Details

B.1 Causal Transition Structure and Parent Set Definitions

  • The causal mask maps state-action tuples to an adjacency matrix; in 2D Navigation, the base local graph selects parent sets (x, ∆x) and (y, ∆y).Source: §B.1 · p. 16

B.2 Parent Distribution

  • The parent sampler fits a GMM to each parent-set marginal and samples sequentially, removing spurious cross-set correlations; samples outside the current local neighborhood are rejected.In the experiments, a single GMM for the sparsest structure was used; HookSweep2 still had some locally out-of-distribution collision-subspace data, much of it involving unreachable overlapping blocks.Source: §B.2 · p. 17

B.3 Dynamics Models

  • The experiments compare unfactored, globally factored, and locally factored dynamics models; the locally factored model uses a separate masked composer and MLP for each child.The authors note that this design does not share potentially useful parent representations across children and may be improved by a shared masked composer.Source: §B.3 · pp. 17–18

B.4 Training Data for the RL Algorithm

The RL training data and algorithm choices varied by experiment. Three offline RL algorithms were implemented, with method-specific details for BCQ and CQL.

B.4 Training Data for the RL Algorithm

  • Generated-data standard deviations were divided by three to avoid samples that were too far out of distribution.The adjustment varied by experiment.Source: §B.4 · p. 18

B.5 Reinforcement Learning Algorithms

  • The implementations added BCQ, CQL, and TD3-BC to Modular RL; TD3-BC followed Fujimoto and Gu.Source: §B.5 · p. 18
  • BCQ used DDPG with a five-component MDN producing 20 action samples per call, trained for 1,000 batches of 2,000, without a perturbation model.Source: §B.5 · p. 18
  • CQL used SAC and balanced a random-policy Q penalty against a behavioral-action Q bonus, with L1 adjustments toward the clipped target-range endpoints.Source: §B.5 · p. 18

C.1 2D Navigation

The appendices describe two continuous-control testbeds where target tasks require generalization beyond the empirical trajectories. They detail each environment’s dynamics, data and augmentation setup, and the HookSweep2 causal-structure heuristic.

C.1 2D Navigation

  • In 2D Navigation, the target is a direct diagonal route from bottom left to top right, while the empirical trajectories support a longer route along their ⌟ shape.The arena uses 2D positions and 2D displacement actions; most regions separate the effects of ∆x and ∆y, but both affect both coordinates in the top-right quadrant.Source: §C.1 · pp. 18–19
  • The 2D Navigation augmentation used parent-set GMMs and KDE-based rejection sampling, then generated 200K-sample datasets, except that EMP repeated the 40K empirical samples.Each GMM had 32 components; non-EMP datasets comprised 40K empirical and 160K generated samples.Source: §C.1 · p. 19

C.2 HookSweep2

  • HookSweep2 requires sweeping two boxes together, although its 1M expert transitions show one box swept while the other remains centered.The 22-dimensional state includes the goal, and actions are 4-dimensional; rewards include small bonuses for keeping the hook near the table.Source: §C.2 · p. 20
  • HookSweep2 used a heuristic causal graph that assumes the hook affects each object at all times, although that relation holds only when they touch.The heuristic treats boxes as separate when more than 5 cm apart; augmentation uses rejection sampling whose density feature map is defined by the box positions.Source: §C.2 · pp. 20–21

C.3 Licenses and Compute

These appendices document compute and licensing, discuss broader-impact concerns, give MOCODA sampling pseudocode, and distinguish its transition sampling from Pearl-style counterfactual reasoning. They also describe assumptions needed for identifiable counterfactual samples and note that those assumptions are restrictive.

C.3 Licenses and Compute

  • Experiments ran on a desktop CPU and NVIDIA GTX 1080 Ti GPU, and the cited code and assets use Apache and MIT licenses.The paper’s implementations were to be released under an open-source license upon acceptance.Source: §C.3 · p. 21

D Further Discussion of Broader Impacts

  • For fairness and robustness applications, causal assumptions merit extra scrutiny because sensitive or socially ascribed attributes may not be directly intervenable and graph choices encode normative assumptions.Source: D Further Discussion of Broader Impacts · p. 21

E MOCODA sampling pseudocode

  • MOCODA pseudocode fits parent-distribution and factored-dynamics models, then samples parent variables and next-state factors to assemble augmented transitions.For overlapping parent sets, it conditions on variables already sampled.Source: E MOCODA sampling pseudocode · p. 22

F Relation to Causal Inference and Counterfactual Reasoning

  • MOCODA samples whole transitions de novo rather than performing Pearl-style counterfactual relabeling; identifiability generally needs further assumptions, and deterministic invertible dynamics are described as restrictive.The models could support subset-variable what-if queries by resampling unspecified parents and affected mechanisms while holding the rest fixed.Source: F Relation to Causal Inference and… · p. 23
Loading 2210.11287v1…