Curious Causality-Seeking Agents Learn Meta Causal World

TL;DR

This work tackles the instability of traditional world models under distributional shifts by introducing a Meta-Causal Graph (MG) that encodes multiple context-specific causal subgraphs activated by latent meta states. A Curious Causality-Seeking Agent actively intervenes to discover meta states and their causal structures, using interventions guided by curiosity rewards (e.g., edge-entropy) and a VQ-VAE-based representation to learn the MG. The authors establish identifiability results for meta states and causal subgraphs under mixed data and interventional conditions, and they validate the approach on synthetic tasks and a robot-arm manipulation task, showing improved prediction accuracy and downstream performance compared to static baselines. This framework advances open-ended world modeling by combining active causal discovery with minimal, robust representations and principled intervention strategies, enabling better generalization across unseen contexts.

Abstract

When building a world model, a common assumption is that the environment has a single, unchanging underlying causal rule, like applying Newton's laws to every situation. In reality, what appears as a drifting causal mechanism is often the manifestation of a fixed underlying mechanism seen through a narrow observational window. This brings about a problem that, when building a world model, even subtle shifts in policy or environment states can alter the very observed causal mechanisms. In this work, we introduce the \textbf{Meta-Causal Graph} as world models, a minimal unified representation that efficiently encodes the transformation rules governing how causal structures shift across different latent world states. A single Meta-Causal Graph is composed of multiple causal subgraphs, each triggered by meta state, which is in the latent state space. Building on this representation, we introduce a \textbf{Causality-Seeking Agent} whose objectives are to (1) identify the meta states that trigger each subgraph, (2) discover the corresponding causal relationships by agent curiosity-driven intervention policy, and (3) iteratively refine the Meta-Causal Graph through ongoing curiosity-driven exploration and agent experiences. Experiments on both synthetic tasks and a challenging robot arm manipulation task demonstrate that our method robustly captures shifts in causal dynamics and generalizes effectively to previously unseen contexts.

Figures (11)

• Figure 1: Illustration of the Meta-Causal Graph concept and the Curious Causality-Seeking Agent framework. Ground Truth: The causal relationship between pushing and opening depends on the latent state (locked vs. unlocked). Limitations of Existing Approaches: Problem 1: Single uniform causal graphs fail to capture context-dependent variations in causal relationships; Problem 2: Domain modeling requires a priori knowledge of state labels, limiting generalization to novel contexts. Our Approach: Curious Causality-Seeking Agents actively intervene to verify causal relationships and discover the critical meta states that determine when causal structures change, enabling the agent to build a comprehensive Meta-Causal Graph without requiring predefined domain labels.
• Figure 2: Performance with different numbers of meta states. (a-b) show prediction accuracy on Chain and Fork environments respectively, while (c-d) show corresponding downstream rewards.
• Figure 3: Comparison of causal patterns discovered with different numbers of meta states. When the number of meta states is set to 2, the model learns two distinct causal patterns (a,d) that correspond to the meta states. With 4 meta states, the model learns four patterns (b,c,e,f), where some patterns share similarities (Patterns 1 and 3), while others occur with lower frequency (Pattern 4) as shown in the frequency distributions (g,h).
• Figure 4: Visualization of environments: (top) Chemical; (bottom) Magnetic.
• Figure 5: Comparison of learned causal graphs in the Chemical task: left—MCG (ours); middle—reference (FCDL); right—ground truth.

Theorems & Definitions (18)

• Definition 1: Intervention graph hauser2012characterization
• Definition 2: Interventional Markov Equivalence hauser2012characterization
• Definition 3: Meta-Causal Graph
• Definition 4: Swap-Label Equivalence
• Theorem 1: Identifiability of Meta States
• Definition 5
• Theorem 2: Identifiability of Overparameterized Meta States
• Proposition 1: Identifiability of Causal Subgraph
• Proposition 2: Effect of Representation Accuracy on Misclassification
• proof