PaperlayerView exact revision on arXiv

Source-linked AI summary

Causal World Models by Unsupervised Deconfounding of Physical Dynamics

Minne Li, Mengyue Yang, Furui Liu, Xu Chen, Zhitang Chen, Jun Wang

arXiv:2012.14228v1cs.LGcs.AI

TL;DR

Conventional world models learn from past sensory regularities without accounting for confounders that affect state transitions, limiting their treatment of intervention counterfactuals. CWMs estimate latent confounders to model alternative futures; the paper reports better performance than state-of-the-art world models across reinforcement-learning and physical-simulation tasks.

  • Problem

    Conventional world models learn spatio-temporal regularities from past data without considering intervention effects or confounders in state transitions.

    Source: §1 · p. 1

  • Method

    CWMs learn abstract object states, estimate latent confounders, and predict alternative futures after object-level interventions.

    Source: §1 · pp. 1–2, §3.3 · p. 4, §3.4 · p. 5

  • Results

    CWMs outperform state-of-the-art world models in reinforcement-learning tasks and physical simulation environments, with reduced sample complexity and improved counterfactual physical reasoning.

    Source: §1 · p. 2, §5 · p. 8

  • Takeaways & Limitations

    Accurate confounder estimates can help model-based reinforcement-learning agents generalize to other parts of the state space and be more robust.

    Source: §1 · p. 1

    The causal-world-model formulation assumes a time-invariant hidden variable blocks all backdoor paths from state to next state, with none from abstract state to observation.

    Source: §3.2 · p. 3, §3.3 · p. 4

Abstract

from arXiv · show

The capability of imagining internally with a mental model of the world is vitally important for human cognition. If a machine intelligent agent can learn a world model to create a "dream" environment, it can then internally ask what-if questions -- simulate the alternative futures that haven't been experienced in the past yet -- and make optimal decisions accordingly. Existing world models are established typically by learning spatio-temporal regularities embedded from the past sensory signal without taking into account confounding factors that influence state transition dynamics. As such, they fail to answer the critical counterfactual questions about "what would have happened" if a certain action policy was taken. In this paper, we propose Causal World Models (CWMs) that allow unsupervised modeling of relationships between the intervened observations and the alternative futures by learning an estimator of the latent confounding factors. We empirically evaluate our method and demonstrate its effectiveness in a variety of physical reasoning environments. Specifically, we show reductions in sample complexity for reinforcement learning tasks and improvements in counterfactual physical reasoning.

1. Introduction

World models aim to let agents imagine alternative futures, but conventional models learn from past observations without accounting for confounders in state transitions. CWMs estimate latent confounders to predict intervention outcomes, and are evaluated for reinforcement-learning sample complexity and counterfactual physical reasoning.

1. Introduction

  • World models seek to let agents internally simulate alternative futures that they have not encountered in the real world.The introduction frames this “dream” capability as a basis for what-if reasoning and decision making.Source: §1 · p. 1
  • Conventional world models learn spatio-temporal regularities from past data but generally do not account for intervention effects or confounding in state transitions.As agents influence the data they gather, the IID assumption may be violated; physical confounders can include inertia, velocity, gravity, and friction.Source: §1 · p. 1
  • CWMs estimate latent confounders and use them to predict alternative futures after object-level interventions to observations.The model learns abstract object states and models their transitions with graph neural networks.Source: §1 · pp. 1–2
  • The paper evaluates CWMs on CoPhy and PHYRE, reporting reduced reinforcement-learning sample complexity and improved counterfactual physical reasoning.The proposed counterfactual risk minimization method is intended to reduce bias from inadequate counterfactual records.Source: §1 · p. 2

2. Related Work

Related work spans causal intervention discovery, visual reasoning, and object-based world models, but much prior work focuses on static settings or requires supervision. CWMs target confounding in state transitions and use unsupervised object abstractions to predict alternative futures.

2. Related Work

  • Causal discovery literature mostly studies static environments, while Lu et al. (2018) address confounding in observations, actions, and rewards rather than state-transition counterfactuals.CWMs instead account for confounders between state transitions to predict alternative futures.Source: §2 · p. 2
  • Prior visual-reasoning and physical-dynamics models often require supervised object or location labels, whereas CWMs extract object abstractions unsupervised.CoPhyNet (Baradel et al., 2020) uses direct object-position supervision and performs counterfactual forecasting in low-dimensional settings.Source: §2 · p. 2
  • Object-based world models often optimize pixel-space prediction or reconstruction, while CWMs learn object-centric abstract states and optimize a hinge loss in latent space.The cited work identifies visually small but informative features as a potential credit-assignment challenge for pixel-space objectives (Łukasz Kaiser et al., 2020).Source: §2 · p. 2
  • Counterfactual learning is used to mitigate bias in historical data, and CWMs set propensity weights from historical sampling strategies rather than the environment state.Related approaches include inverse propensity score methods for unbiased learning.Source: §2 · pp. 2–3

3. Causal World Models

Causal World Models extend conventional world models to predict alternative observation trajectories under interventions by estimating latent, time-invariant confounders. The framework combines object-level state abstractions, causal identification, and doubly robust learning from historical observations.

3.2. Causal POMDPs

  • Unlike conventional POMDP world models, causal POMDPs account for hidden, time-invariant factors that influence both the intervened state and its outcome.The paper describes physical confounders such as object mass, friction, and gravitational forces; observational and interventional conditionals can differ because of them.Source: §3.2 · pp. 3–4, §3.1 · p. 3

3.3. Learning Causal World Models

  • CWM predicts dream-world trajectories after an object-level observable intervention, using the altered observation and an estimate of latent confounders.The target is the alternative future following a changed initial observation; training uses offline experience and can evaluate a given dream policy.Source: §3.3 · p. 4
  • Under Assumption 1, an estimator of latent confounders makes the intervention distributions for current and future observations identifiable.The proof applies the back-door criterion to identify those distributions from observed data, intervened observations, and the confounder estimator.Source: §3.3 · p. 4

3.4. Model Design

  • CWM encodes each observation into object-level abstract states and uses a graph neural network to model pairwise interactions and predict state updates.The encoder uses K object slots; the transition model adds predicted updates to the abstract states.Source: §3.4 · p. 5
  • CWM trains its transition model with an energy-based hinge loss that contrasts positive transitions with corrupted negative samples.The objective is averaged over samples from the experience buffer, with a margin hyperparameter.Source: §3.4 · p. 5
  • A recurrent estimator approximates each object's latent confounders from its state trajectory and supplies those estimates to the transition model.The estimator uses a shared-parameter GRU across object trajectories and concatenates its estimate with the state before transition prediction.Source: §3.4 · pp. 5–6

3.5. Doubly Robust Learning from Historical Observations

  • The doubly robust oracle is unbiased against the true trajectory observation given the propensity score, and its objective is used to train CWM.The propensity score represents the historical observation distribution, addressing counterfactual risk under continuous interventions and historical sampling policies.Source: §3.5 · p. 6

4. Experiment

On CoPhy, CWMs and CRM-CWMs lead the reported dream-quality comparisons; on PHYRE, CWMs improve task solving relative to WMs and reduce sample complexity relative to model-free agents.

4. Experiment

  • The study compares causal world models with conventional world models on CoPhy for dream-creation quality and PHYRE for dream-world usability, reporting reduced reinforcement-learning sample complexity and improved counterfactual physical-reasoning predictions.CoPhy uses 700 training and 300 test tuples of two trajectories, each with 29 time steps; counterfactual prediction is evaluated using H@1 and MRR.Source: §4 · p. 6, §4.1 · p. 6, §4.2 · p. 7

4.1. Dream Quality

  • CRM-CWMs estimate a density of historical training samples for propensity weighting, while the compared models share the object extractor, encoder, and transition estimator.Models are reported using their best configuration across slot counts from two to six.Source: §4.1 · pp. 6–7
  • CWMs and CRM-CWMs consistently achieve the best H@1 and MRR results across training environments, while qualitative examples show stronger counterfactual trajectories than WMs.The analysis attributes conventional WMs' weaker results to sensitivity to environment difficulty and hidden confounders.Source: §4.1 · p. 7

4.2. Dream Usability

  • On PHYRE, agents rank candidate object placements by predicted task success, then attempt placements on previously unseen tasks.The task uses a default set of 100,000 action candidates, and agents receive a binary reward after rollout.Source: §4.2 · pp. 7–8
  • On PHYRE, CWMs outperform WMs in sample complexity and final average success within 50 attempts, while model-based agents reduce sample complexity versus model-free agents.Success-percentage curves are averaged over all test tasks and three random seeds.Source: §4.2 · p. 8

5. Conclusion

The paper presents CWMs as an unsupervised way to predict alternative futures by estimating latent confounders, and reports stronger performance than state-of-the-art WMs in reinforcement learning and physical simulation. It also uses Simpson’s paradox to illustrate how confounding can reverse an observed trend when groups are combined.

5. Conclusion

  • CWMs estimate latent confounders to predict alternative futures without supervision, outperforming state-of-the-art WMs in reinforcement learning and physical simulation.The paper also proposes counterfactual risk minimization to address bias in counterfactual datasets and reports its effectiveness for learning physical dynamics.Source: §5 · p. 8
  • The authors identify broader applications and more complex environments as future directions for CWM evaluation.Source: §5 · p. 8

A. Simpson’s Paradox in the Medical Treatment Scenario

  • In the treatment example, overall recovery rises from 40% to 50% while rates fall for both male and female patients, illustrating a trend reversal associated with confounding.The passage attributes the reversal to gender affecting both treatment and recovery.Source: A. Simpson’s Paradox in the Medical… · p. 14

B. Proof of Proposition 1

The proof introduces notation for a trajectory’s latent state.

B. Proof of Proposition 1

C. Evaluation Metrics for the Dream Quality Experiment

Dream-quality performance is evaluated in latent space by ranking predicted abstract states against the encoded ground truth and reference states.

C. Evaluation Metrics for the Dream Quality Experiment

D.1. Dream Quality Experiments

The dream-quality experiments train conventional, causal, and CRM-augmented world models on CoPhy trajectories, with separate model configurations for each variant.

D.1. Dream Quality Experiments

  • Training uses 700 episodes of 29 steps and tests on 300 episodes of 29 steps, with RGB observations resized to 50 × 50 pixels.Models train for 200 epochs with Adam at 5e −4 and batch size 25.Source: §D.1 · p. 15

D.1.1. WMS

  • WMs encode each object as a 100-dimensional vector, and their node and edge transition models use MLPs matching the object encoder architecture.The encoder uses two 512-unit ReLU layers, with LayerNorm after the second, followed by a 4-unit layer.Source: §D.1.1 · p. 15
  • The WM hinge-loss margin is γ = 1, and its squared-distance term is scaled by 0.5/σ2 with σ = 0.5.Source: §D.1.1 · p. 15

D.1.2. CWMS

  • CWMs estimate confounders with a two-layer GRU of hidden dimension 32 and concatenate that estimate with the abstract state before transition prediction.Source: §D.1.2 · p. 15

D.1.3. CRM-CWMS

  • CRM-CWMs estimate propensity scores with a Gaussian density function and use Equation 5 as their training objective.Other settings match those of CWMs.Source: §D.1.3 · p. 15

D.2. Dream Usability Experiments

The dream-usability experiments describe the PHYRE task setup and the DQN and world-model components used by the agents.

D.2. Dream Usability Experiments

  • PHYRE experiments use 70 training tasks and 30 test tasks, with agents ranking the same 200 actions per test task.Observations are 256 × 256 images encoded as seven-channel inputs.Source: §D.2 · p. 16

D.2.1. DQN

  • The DQN combines encoded actions and observations in a fusion module that predicts reward, and it is trained end-to-end with cross-entropy.The action is three-dimensional, and DQN uses Adam with learning rate 3e −4.Source: §D.2.1 · p. 16

D.2.2. WMS

  • The WMS setup uses an object extractor with K typically set to the task's object count, such as K = 3 for Figure 5a, and a separate state-status classifier.The classifier predicts solved versus unsolved from an abstract state and is trained with cross-entropy.Source: §D.2.2 · pp. 16–17
  • In the 2-ball CoPhy illustration, predicted dream trajectories for WMs, CWMs, and CRM-CWMs are compared with ground truth using same-scale PCA projections and episode-level MSE.Source: Figure 2 · p. 17
  • All models train for 200 epochs with Adam at 5e −4 and batch size 16; experiments finish on one GTX 1080 Ti GPU in under six hours.Source: §D.2.2 · p. 17

D.2.3. CWMS

  • The CWMS setup follows the conventional WMS settings while concatenating the estimated confounder with the abstract state before transition prediction.Source: §D.2.3 · p. 17

E.1. Qualitative Examples of Dream Quality Experiments

The appendix provides additional counterfactual dream-trajectory examples for CoPhy environments with three and four balls.

E.1. Qualitative Examples of Dream Quality Experiments

  • Additional dream trajectories are provided for environments shown in Figures 2 through 6.Source: §E.1 · p. 17
  • Figure 3 compares WM, CWM, and CRM-CWM dream trajectories with ground truth in the 3-ball environment using same-scale PCA projections and episode-level MSE.Source: Figure 3 · p. 18
  • Figure 4 presents the corresponding trajectory comparison in the 4-ball environment, with projected paths and episode-level MSE.Source: Figure 4 · p. 18

E.2. Ablation Study on the Choice of K in Dream Quality Experiments

The appendix reports slot-count tuning results and illustrates dream trajectories in the five- and six-ball CoPhy environments.

E.2. Ablation Study on the Choice of K in Dream Quality Experiments

  • CWMs and CRM-CWMs benefit from larger slot counts K across environments, suggesting larger K as a choice when a new environment lacks prior information.The main-paper plots use each model's best K configuration.Source: §E.2 · p. 18
  • Figure 5 compares WM, CWM, and CRM-CWM dream trajectories with ground truth in the 5-ball environment using same-scale PCA projections and episode MSE.Source: Figure 5 · p. 19
  • Figure 6 presents the corresponding trajectory comparison in the 6-ball environment, with same-scale PCA projections and episode-level MSE.Source: Figure 6 · p. 19
  • Figure 7 presents H@1 ranking results for multi-step latent-space prediction across slot-count variations.Source: Figure 7 · p. 20
  • Figure 8 presents MRR ranking results across slot-count variations and reports that Causal SWM achieves the best result consistently.Source: Figure 8 · p. 21

Related papers

Loading 2012.14228v1…