PaperlayerView exact revision on arXiv

Source-linked AI summary

Counterfactual Data Augmentation using Locally Factored Dynamics

Silviu Pitis, Elliot Creager, Animesh Garg

arXiv:2007.02863v2cs.LGcs.AIcs.ROstat.ML

TL;DR

Interacting processes often have sparse local interactions, although their global causal models can be densely connected. The paper introduces local causal models and CoDA, which stitches locally independent observed subprocesses into counterfactual training data. CoDA significantly improves reinforcement-learning performance across locally factored tasks, including batch-constrained and goal-conditioned settings.

  • Problem

    Interacting subprocesses often have sparse interactions, but global causal models can be densely connected, limiting the usefulness of global factorization.

    Source: §1 · p. 1

  • Method

    The paper conditions global causal models on subsets of the state space and uses CoDA to swap locally independent components from observed transitions into counterfactual training samples.

    Source: §1 · p. 2

  • Results

    CoDA significantly improves sample efficiency in standard, batch-constrained, and goal-conditioned reinforcement-learning settings.

    Source: §1 · p. 2

  • Takeaways & Limitations

    CoDA can more than double sample efficiency and final performance for reinforcement-learning agents in locally factored environments.

    Source: §6 · p. 9

    The reported approach uses fully observable, disentangled state spaces with a fixed top-level decomposition.

    Source: §6 · p. 9

Abstract

from arXiv · show

Many dynamic processes, including common scenarios in robotic control and reinforcement learning (RL), involve a set of interacting subprocesses. Though the subprocesses are not independent, their interactions are often sparse, and the dynamics at any given time step can often be decomposed into locally independent causal mechanisms. Such local causal structures can be leveraged to improve the sample efficiency of sequence prediction and off-policy reinforcement learning. We formalize this by introducing local causal models (LCMs), which are induced from a global causal model by conditioning on a subset of the state space. We propose an approach to inferring these structures given an object-oriented state representation, as well as a novel algorithm for Counterfactual Data Augmentation (CoDA). CoDA uses local structures and an experience replay to generate counterfactual experiences that are causally valid in the global model. We find that CoDA significantly improves the performance of RL agents in locally factored tasks, including the batch-constrained and goal-conditioned settings.

1 Introduction

Interacting subprocesses often have sparse, time-local interactions, making global factorization impractical despite local causal independence. The paper formalizes local causal models and uses them in CoDA to augment reinforcement-learning data with causally valid counterfactuals.

1 Introduction

  • Locally independent subprocesses can be leveraged even when they interact globally, unlike a globally factored model that assumes they never interact.The motivating observation is that subprocesses are causally independent between interactions.Source: §1 · p. 1
  • Of three illustrated proposals, only the second is accepted: its factual sources and resulting sample meet the local-factorization requirements.The accepted sample can be used as additional reinforcement-learning training data.Source: Figure 1 · p. 2
  • CoDA swaps factorized subspaces from observed trajectory pairs to create counterfactually plausible samples under mild conditions, without requiring a forward dynamics model.Because it modifies training data, CoDA is compatible with any agent architecture, including unfactored agents.Source: §1 · p. 2
  • CoDA significantly improves sample efficiency in standard, batch-constrained, and goal-conditioned reinforcement learning settings.The paper reports these results using an attention-based method to discover local causal structure in disentangled state spaces.Source: §1 · p. 2
  • The paper defines local causal models conditioned on a subset of the state space and presents CoDA as a generalized augmentation strategy that includes goal relabeling and visual augmentation as instances.It also proposes locally conditioned CoDA, which swaps independent subprocesses to form counterfactual experiences.Source: §1 · p. 2

2 Local Causality in MDPs

The section sets up factored Markov decision processes and explains why global causal graphs can remain dense even when interactions are rarely active. Local causal models condition on a state-action subspace where dependencies disappear, enabling simpler counterfactual reasoning.

2.1 Preliminaries and Problem Setup

  • The setup assumes a known, fixed top-level decomposition of the state space; hierarchical decomposition, multiple representations, dynamic factor counts, and learned latent representations are left to future work.The action space may also be decomposed.Source: §2.1 · pp. 2–3
  • A factored MDP represents each next-state component through its parents and noise, with independent noise variables and a structurally minimal causal graph.Under structural minimality, a parent is a variable whose value can directly affect its child through the structural equation.Source: §2.1 · p. 3

2.2 Local Causal Models (LCMs)

  • In most practical scenarios the global causal DAG is fully or nearly connected, so global factorization rarely offers an advantage over a monolithic model.A possible interaction in even one state can preserve an edge in the global graph, despite that interaction being rare.Source: §2.2 · p. 3
  • An LCM conditions on a subspace of states and actions where dependencies may disappear, yielding a sparser local DAG and restricted structural equations.For two robot arms, conditioning on them being far apart yields two separate components.Source: §2.2 · pp. 3–4, Figure 2 · p. 4
  • When an intervention stays within the local model’s conditioning subspace, counterfactual reasoning can use that submodel and consider fewer structural equations.In the robot example, reasoning about each arm can be done separately when the arms are far apart.Source: §2.2 · p. 4

3 Counterfactual Data Augmentation

CoDA generates counterfactual transitions by swapping locally independent components between observed transitions. It validates proposed swaps against local causal structure and can infer that structure from data.

3 Counterfactual Data Augmentation

  • Goal relabeling and visual feature augmentation use global independence, while CoDA also swaps components that are independent only within local causal structures.Goal relabeling may use the known reward mechanism to relabel rewards; visual augmentation changes features treated as independent of physical dynamics.Source: §3 · p. 5
  • Transitions sharing a local factorization provide unbiased samples of independent mechanisms that can be swapped to create counterfactual data, provided the result remains in the local subspace.The swapped components come from observed transitions, rather than requiring a forward dynamics model.Source: §3 · pp. 4–5
  • n^m CoDA samples can be formed from n independent samples across m connected components, an exponential increase in combinations.The count assumes the samples are independent and the subspace has m connected components.Source: §3 · p. 5
  • CoDA can mix locally independent components across multiple time steps, beyond its usual single-transition setting.The paper describes single-transition augmentation as typical for RL agents, while multi-step mixing is also possible.Source: §3 · p. 5
  • Counterfactual data shifts the training distribution, and restricting modifications to local subspaces biases replay toward more factorized transitions.The paper uses heuristic observed-to-counterfactual ratios and leaves full characterization of CoDA’s selection bias to future work.Source: §3 · pp. 5–6
  • CoDA masks each transition, finds independent components, swaps one between transitions, then accepts the proposal only if its mask preserves the same graph partitions.For swaps across neighborhoods, the mechanisms must also share structural equations.Source: §3 · p. 6
  • Local factorization can be inferred with a state- and action-conditioned network mask; the paper uses a set transformer to infer structure, not to generate future states.This data-driven route assumes learning local factorization is easier than modeling the dynamics.Source: §3 · pp. 6–7

4 Experiments

Across online, batch, and goal-conditioned settings, CoDA improved performance over non-CoDA baselines. Experiments include learned and ground-truth masks, batch data expansion, and augmenting HER in two goal-conditioned tasks.

4 Experiments

  • CoDA significantly improves agent performance over non-CoDA baselines in online, batch, and goal-conditioned RL.The experiments assess CoDA with different agents and data settings.Source: §4 · p. 7
  • Ground-truth- and learned-mask CoDA improve TD3 sample complexity in the bouncing-ball tasks, while Dyna performs poorly from model bias.The replay buffer is expanded about 8×; the learned mask is inferred using a transformer pretrained on approximately 42,000 random-policy samples.Source: §4 · p. 7, Figure 4 · p. 7
  • In Pong batch RL, learned-mask CoDA more than doubles effective data size and yields a 3× performance boost at smaller data sizes.A 1R:5C Real:CoDA ratio performs slightly worse than 1R:3C, attributed to distributional shift.Source: Table 1 · p. 8, §4 · p. 8
  • In FetchPush and Slide2, HER with a CoDA-enlarged dataset approximately doubles sample efficiency compared with base HER.The heuristic mask treats objects more than 10 cm apart as disentangled and works in both environments.Source: Figure 5 · p. 8, §4 · p. 8

5 Related Work

Related work connects CoDA to factored dynamics, causal counterfactual methods, and data augmentation. These approaches differ in what structure they exploit and the scale at which they operate.

5 Related Work

  • Factored MDPs assume each state variable has a fixed subset of parent variables, while context-specific independence compactly represents conditional structure.The cited CSI methods require discrete variables, whereas CoDA is applied to continuous tasks.Source: §5 · p. 8
  • Prior causal RL work uses counterfactual trajectories to reduce off-policy gradient variance, whereas CoDA factors representations within a single transition.The cited Buesing et al. approach uses coarse-grained trajectory-level representations.Source: §5 · p. 9
  • Supervised-learning augmentation encodes causal invariance to input perturbations, and counterfactual augmentation is also discussed in connection with robustness and debiasing.The passage presents robust performance and debiasing as relevant where fairness is a concern.Source: §5 · p. 9

6 Conclusion

The paper introduces local causal models to recover useful decomposition when global causal structure is densely connected, and uses them for counterfactual data augmentation. It reports more than doubled sample efficiency and final performance in locally factored RL environments, while identifying broader settings for future work.

6 Conclusion

  • CoDA can more than double RL agents’ sample efficiency and final performance in locally factored environments.The conclusion also points to future work on learned masks, prioritization, model-based reasoning, partially observable or entangled settings, and applications beyond RL.Source: §6 · p. 9
  • The LCM framework captures decomposition benefits even when the global causal model is densely connected.The authors identify fully observable, disentangled states with a fixed top-level decomposition as a current scope boundary, with partially observable and entangled settings left for future work.Source: §6 · p. 9

Broader Impact

The broader impact discussion weighs CoDA’s potential to support robustness and fairness against risks from distribution shift, counterfactual errors, and empowered agents. It also notes that batch-constrained use may not remedy biased or missing data without further work.

Distribution shift

  • CoDA-generated distribution shift may create safety risks when modules trained only on observational data perform poorly on counterfactual data.The authors anticipate uncertainty methods will be important for controlling these risks.Source: Broader Impact · p. 10

Potential benefits

  • Counterfactual out-of-distribution training could support robustness and fairness, including potential fairness-focused applications in reinforcement learning and computer vision.The paper cites counterfactual gender-pronoun swaps in NLP as an example of fairness-oriented augmentation.Source: Broader Impact · p. 10

Batch-constrained settings

  • Batch-constrained CoDA can expand modest observational datasets, but subspace swapping alone may not produce plausible exploratory data or correct sample bias.The authors call for applied work characterizing how sample bias affects CoDA; in RL, distribution shifts also require off-policy learning or high-variance corrections for on-policy algorithms.Source: Broader Impact · p. 10

Risks

  • A poorly fit or non-generalizing local factorization model could yield invalid counterfactual data or poorly formulated imagined goals, hurting agent performance.The discussion also flags risks from misguided goals, poor modeling, resource constraints, and inadequate representation of human values.Source: Broader Impact · p. 10

A Proof of Proposition 1

The proof establishes when causal mechanisms remain independent across a union of local domains and shows that conditioning on a subset can yield a sparser graph. The accompanying online experiments report a CoDA advantage in sparse-reward tasks, except Place 4, where no agent earns reward.

A Proof of Proposition 1

  • When L is a subset of X, its local causal graph is sparser than the graph for X.The proof applies minimality in X to any parent present in the local graph, so that parent must also appear in the graph for X.Source: A Proof of Proposition 1 · p. 16
  • Two mechanisms are independent in the union of local domains if and only if they are independent in each domain and their mechanisms agree.The proof’s reverse direction rules out a union-only edge by comparing states from the different domains and using agreement of the mechanism on shared parent values.Source: A Proof of Proposition 1 · p. 16

B Additional Experiment Details

  • The online experiment keeps the TD3 architecture and hyperparameters largely fixed while modifying the agent’s training distribution.The appendix describes Spriteworld extensions and CoDA sampling from replay transitions, using either a ground-truth or learned mask.Source: §B.1 · p. 16

B.1 Online RL

  • CoDA agents maintain a clear advantage in four sparse-reward Spriteworld task variants, except Place 4, where no agent achieves reward.Reward is granted only when all sprites reach their targets; the supplied figure reports average reward and likewise notes the Place 4 exception.Source: §B.1 · p. 17, Figure 6 · p. 17

B.2 Batch RL

The batch-RL appendix describes dataset expansion configurations for Pong and the training setup, but the supplied passages do not include the table’s result values. Its configuration notes contrast the harmful CoDA forward model with an MBPO-based Dyna variant.

Implementation

  • The experiment builds a dataset from real, Dyna, and/or CoDA data before filling the TD3 replay buffer and running 500,000 optimization steps.The buffer is expanded to hold the entire enlarged dataset.Source: §B.2 · p. 17

Environment

  • The Pong setup uses a 12-dimensional state and sparse ±1 rewards when the ball is behind a player’s paddle.The environment is adapted from RoboSchoolPong-v1, with episode resets after scoring or 150 steps.Source: §B.2 · p. 18

Configurations

  • Using the CoDA model as a Dyna forward model harms performance, while five-step MBPO rollouts do better than one-step rollouts.The authors attribute the MBPO result to lower model bias and a more diverse dataset; the one-step result is not shown.Source: §B.2 · p. 18
  • The MBPO + CoDA configuration expands the original dataset to five times its size: one part real, one part MBPO, and three parts CoDA.The described sequence first doubles the dataset with MBPO, then adds CoDA data to reach the 5x total.Source: §B.2 · p. 19

B.3 Goal-conditioned RL

In the goal-conditioned experiments, CoDA outperforms the base HER agent on all tested settings, using a hand-coded distance heuristic to identify local independence. The section also describes SANDy models for inferring local factorization and evaluates them in synthetic processes and Spriteworld.

B.3 Goal-conditioned RL

  • CoDA outperforms the base HER agent on all tested goal-conditioned settings.The experiments use default CoDA hyperparameters; FetchPush-v1 drops state features that entangle object and gripper positions, while Slide2 uses sparse rewards for placing both pucks at their ordered targets.Source: §B.3 · p. 20
  • A hand-coded CoDA heuristic treats actions as entangled with the gripper and separates objects from the gripper/action when they are more than 10 cm away.The authors describe the heuristic as generous and suboptimal; it prioritizes high precision over recall, and they note that a learned mask could also use velocity.Source: §B.3 · p. 20

C Inferring Local Factorization

  • SANDy infers local factorization from next-state prediction, with ground-truth connectivity used to evaluate—not train—the models.SANDy-Transformer is used for the paper’s online RL experiments.Source: C Inferring Local Factorization · p. 21

C.1 Methods

  • SANDy-Mixture thresholds a dynamics network’s Jacobian to estimate local dependencies, while SANDy-Transformer uses self-attention across embedded state and action components.The Transformer predicts next-state features by discarding action outputs and is trained with mean squared error; the mixture uses sparse local dependencies and an attention mechanism.Source: §C.1 · pp. 21–22

C.2 Evaluation environments

  • The synthetic evaluation contrasts globally factorized dynamics with a variant that adds global transitions when a local factor’s norm exceeds ϵ.Lower ϵ produces more global interaction; the evaluation also includes Spriteworld with a ground-truth mask renderer.Source: §C.2 · p. 23

C.3 Results

SANDy-Mixture performs well on simpler synthetic Markov processes but scales poorly to Spriteworld, where SANDy-Transformer performs favorably. The section also reports that autoregressive dynamics-model rollouts fail to capture collisions and long-term dependencies.

C.3 Results

  • Connectivity labels are available for evaluation in these controlled environments but are not used to train SANDy; a real-world evaluation protocol is left to future work.The models are evaluated on held-out transitions by sweeping the mask threshold and measuring true and false positive rates.Source: §C.3 · p. 23
  • SANDy-Mixture reaches average AUCs of 0.96 on stationary MP and 0.91 on non-stationary MP, but scales poorly to Spriteworld, where SANDy-Transformer performs favorably.The ROC evaluation uses held-out connectivity labels that are not used to train the attention model.Source: Figure 8 · p. 24, §C.3 · p. 24

D Fitting dynamics models to Spriteworld

  • Autoregressive rollouts from Linear, MLP, and LSTM dynamics models fail to capture collisions and long-term dependencies despite relatively low next-state prediction error.The models were trained to predict the next state from the current state and action; the LSTM predicts successive next-state dimensions rather than multiple time steps.Source: D Fitting dynamics models to Spriteworld · p. 24

E Sample efficient dynamics modeling with CoDA

The experiments test whether counterfactual CoDA transitions improve forward-dynamics model training in Spriteworld. Adding unique CoDA transitions improves validation loss and prevents overfitting, while autoregressive dynamics-model rollouts remain a weaker alternative for generating trajectories.

E Sample efficient dynamics modeling with CoDA

  • Adding 35,000 unique counterfactual CoDA transitions significantly improves validation loss and prevents overfitting compared with training from 2,000 random-policy transitions.The baseline and an identity-mask random CoDA strategy overfit; random CoDA does somewhat better than no CoDA, which the authors attribute to regularization.Source: E Sample efficient dynamics modeling… · pp. 24–25
  • With 2,000 base transitions, the authors estimate that as many as 6.5 trillion counterfactual samples could be generated under their stated collision-free and connected-component assumptions.Their estimate assumes 80% of transitions involve no collisions and four connected components per transition.Source: E Sample efficient dynamics modeling… · p. 25
  • Autoregressive rollouts from fitted dynamics models can have low next-state error yet fail to capture collisions and long-term dependencies.The sampled trajectories share the same initial state.Source: Figure 10 · p. 26

F Compute Infrastructure

Experiments ran on local machines and a compute cluster using several GPU types; the authors report that all experiments can run locally.

F Compute Infrastructure

  • Experiments used GTX 1080 Ti, Titan XP, and Tesla P100 GPUs for parallel jobs, but can all run locally, with GPUs optional for Spriteworld and Pong and recommended for Fetch.Source: F Compute Infrastructure · p. 26

Related papers

Loading 2007.02863v2…