How Causal Inductive Biases Shape Structured World Models: VCD, COMET & SPARTAN

Introduction

World models are central to reinforcement learning and agent research: if an agent can learn an environment prediction model, it can simulate future trajectories for planning. Recent world models excel in games, robotics, and visual control, but most use dense, monolithic neural networks that compress dynamics into an opaque black box. This thesis asks: since the real world consists of sparse, local, reusable causal mechanisms, should world models explicitly exploit that structure?

The author, Anson Lei (Mansfield College, University of Oxford, 2025), organizes the thesis around three progressive questions: how to learn causal representations, how to learn independent mechanisms, and how to discover local dependency structures. Three corresponding methods are proposed: Variational Causal Dynamics (VCD), COMET (competition-of-experts for independent mechanisms), and SPARTAN (sparse Transformers for local dependency discovery). The work demonstrates that structured world models not only predict accurately but also adapt modularly, resist spurious correlations, and provide interpretable internal structure. Thesis link: https://ora.ox.ac.uk/objects/uuid:e14fe66f-83bf-4faa-a658-b87e3a168f55/files/dj38607754

Background: From Representations to Mechanisms

Early world models compress observations into a continuous latent space and predict future latents with a recurrent or feedforward dynamics network. This improves computational efficiency but latent variables may not correspond to interpretable real-world factors. Object-centric representations decompose scenes into slots, yet still lack explicit modeling of object interactions. The thesis frames the problem in a causal modeling framework: each variable in a causal graph is determined by its parents and a mechanism; environment changes are interventions on a few mechanisms. The Independent Causal Mechanisms principle and the Sparse Mechanism Shift hypothesis suggest that if a world model explicitly learns sparse causal structure, it can adapt to new environments by updating only the affected mechanisms.

Three structural biases are reviewed: causal representation learning (latent variables correspond to true factors, dynamics graph is sparse), independent mechanism learning (dynamics decomposed into reusable interaction primitives), and structure discovery (identifying local dependencies among entities). The subsequent chapters each address one bias.

Variational Causal Dynamics (VCD)

Method

VCD jointly learns latent representations and a sparse causal dynamics graph within a latent state-space model. Latent variables are treated as causal variables, each depending only on its causal parents. The dynamics model comprises independent mechanisms, sparse causal dependencies, and sparse interventions. Different environments share unperturbed mechanisms; only intervened variables receive environment-specific parameters. When entering a new environment, the model identifies which mechanisms changed and updates only those parts, avoiding full retraining.

Experiments

Evaluated on a 2D multi-body simulation with particles interacting via spring-like or electromagnetic forces, under both mixed-state and image observations. Baselines: RSSM (monolithic) and MultiRSSM (mixture of experts). VCD achieves lower prediction error than RSSM and matches or approaches MultiRSSM. Crucially, VCD learns more disentangled representations, accurately identifying sparse causal relations and intervention targets. Modified MCC scores: 0.975 (mixed-state) and 0.908 (image), both significantly higher than RSSM.

Key Conclusion

Sparse causal graphs serve as an inductive bias that not only improves interpretability but also drives disentangled representation learning and modular adaptation.

Learning Independent Mechanisms via Competition (COMET)

Motivation

Humans understand the world through reusable rules: collision, stacking, rolling, attraction, repulsion, speed limits. These rules recombine across environments. COMET learns a set of independently parameterized mechanism modules and composes them in new environments.

Two-Stage Training

Stage 1 (Competition): Multiple mechanisms explain data from multiple environments; the mechanism with lowest prediction error receives the gradient update. Stage 2 (Composition): In a new environment, the model learns how to select and combine the previously learned mechanisms.

Results and Challenges

COMET learns identifiable mechanisms in image-based environments and adapts to new environments with varying object counts using fewer samples. However, competition can suffer from mode collapse: if some mechanisms dominate early, others fail to learn. Warm-start strategies and mechanism validity windows affect training stability.

Contribution

Shows that structured world models need not only variable disentanglement but also rule disentanglement. When environment changes correspond to mechanism recombination, mechanism-level modularity is more natural than whole-parameter updates.

Learning Local Dependency Structures with Sparse Transformers (SPARTAN)

From Global to Local Graphs

Fixed global causal graphs are inflexible for complex physical scenes: not all objects interact at all times. Interactions are local and state-dependent. SPARTAN uses sparse hard attention to learn object-level local dependencies. The model takes object embeddings as input; sparse attention selects the true causal parents for the current state, turning Transformer attention patterns into interpretable local causal graphs.

Robustness and Few-Shot Adaptation

Evaluated on Interventional Pong, CREATE, and Traffic environments. Compared to standard Transformers and global graph baselines, SPARTAN maintains predictive performance while more accurately recovering local causal graphs and adapting better with few trajectories. Robustness tested by removing non-causal entities: if a model relies on spurious correlations, removing irrelevant objects degrades predictions; if it captures true causal parents, predictions remain stable. SPARTAN resists such distractors better than unregularized Transformers.

Significance

SPARTAN combines Transformer expressivity with causal structure discovery: attention is shaped by sparsity constraints into an interpretable local dependency graph, not merely a visualization tool.

Sparse Transformers Beyond World Models

Direct Advantage Estimation in RL

SPARTAN is applied to learn value and advantage functions directly from image patches. Sparse attention identifies which image regions influence decisions. Benefits: explicit advantage learning improves policy generalization; sparse attention reveals whether the policy attends to task-relevant objects rather than extrinsic rewards or background distractors.

Language Model Mechanistic Interpretability

Standard dense attention in Transformers is replaced with SPARTAN-style sparse hard attention as a post-training technique on GPT-2 and LLaMA-1B. The sparse models explain a comparable proportion of behavioral effects with fewer attention heads. This suggests sparse attention can help build more interpretable language model internals, though it remains preliminary.

Conclusion

Main Contributions

The thesis argues that when a model's computational structure respects the world's causal structure, adaptability, robustness, and interpretability naturally improve. VCD demonstrates sparse causal graphs for representation disentanglement and modular adaptation; COMET shows competition-based learning of reusable independent mechanisms; SPARTAN uses sparse attention to discover local, state-dependent causal structures. The three are complementary: VCD focuses on causal representation learning, COMET on mechanism learning, SPARTAN on local structure discovery. Together they show structured world models should pursue not only low prediction error but also explainability, transferability, and fast adjustment in changing environments.

Limitations and Future Work

Experiments rely on simulated environments with controlled interventions; real-world observation noise, partial observability, and complex interactions are harder.

Causal identifiability depends on data distribution and inductive biases; if true mechanisms are highly coupled or graphs dense, sparsity assumptions may fail.

Scaling SPARTAN and similar sparse structures to larger models and real robot tasks needs further validation.

Future directions include combining the three approaches: integrating SPARTAN's sparse attention regularization with VCD's joint representation learning, introducing COMET's mechanism competition into Transformer-like architectures, and testing structured world models in robotics, autonomous driving, and long-horizon agent tasks.

Final Perspective

The thesis reminds us that the goal of world models is not merely "predicting future pixels" but learning how the world organizes itself. A truly useful world model should know which variables matter, which mechanisms are reusable, which changes require adaptation, and which disturbances can be ignored. Causal inductive biases are a step toward that goal. For agent research, this line is critical: future autonomous systems must act continuously in changing environments; if their internal models are dense black boxes, diagnosing failures, adapting quickly, and deploying safely become difficult. Structured world models offer a more controllable alternative: models that not only predict the world but also understand it in a way that mirrors the world's own structure.

Code example

来源：专知
本文
约4000字
，建议阅读
7
分钟
这篇 Oxford 博士论文提出一个清晰问题：真实世界往往由稀疏、局部、可复用的因果机制组成，那么世界模型是否也应该显式利用这种结构？