Curious Causality-Seeking Agents in Open-ended Worlds
Authors: Zhiyu Zhao
Haoxuan Li
Haifeng Zhang
Jun Wang
Francesco Faccio
Jürgen Schmidhuber
Mengyue Yang
Venue: NeurIPS 2025 poster ()
Keywords: Causality
World Model
Primary area: probabilistic_methods

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. However, in truly open-ended environments, the apparent causal mechanism may drift over time because the agent continually encounters novel contexts and operates within a limited 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 Meta-Causal Graph as world models for open-ended environments, 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 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.

Curious Causality-Seeking Agents in Open-ended Worlds
Zhiyu Zhao', Haoxuan Li², Haifeng Zhang Jun Wang", Francesco Faccio
Jürgen Schmidhuber Mengyue Yang"*
University of Bristol PPeking University
Institute of Automation, Chinese Academy of Sciences
Schooo of Artificial Intelligence, Chinese Academy of Sciences
University College London
King Abdullah University of Science and Technology The Swiss AI Lab, IDSIA-USI/SURI
Abstract
When building a world model, a common assumption is that the environment has a single, unchanging underlying causal rule. like applying Newton's 's laws every situation. However, in truly open-ended environments, the apparent causal mechanism may drift over time because the agent continually encounters novel contexts and operates within a limited observational window. This brings about problem that, when building a world model, even subtle shifts in policy environment states can alter the very observed causal mechanisms. In this work, we introduce the Meta-Causal Graph world models for open-ended environments, 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, Agent whose objectives are to (1) identify
we introduce a Causality-Seeking 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 challenging robot arm manipulation task demonstrate that our method robustly captures shifts causal dynamics and generalizes effectively to previously unseen contexts.
Introduction
World models 21 70 have emerged critical component in reinforcement learning, enabling agents to simulate and plan in complex environments. These models aim to capture the underly- ing dynamics of the environment, allowing agents to predict future states and evaluate potential However, simply learning which variables correlate can mislead an agent when the world's 's dynamics change. Recently, causality has been adopted to improve world models for interactive and complex environments, since causal descrrinng how one factor brings about another another-capture capture the underlying data- generatierrnn process and yield more robust and generalizable decision-making [36 59]. While many existing approaches incorporate notions of causality, they typically fail to capture the open-world causal relationships.
actions 23 21 6 55
To identify the underlying world rules, traditional causal models rely predominantly on observational data [46), implicitly assuming fixed causal structure. In open-ended environments, although the
CCorresponiing author: mengyue. yang @ abristol.ac.uk
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Push
Push
Open
Intervention
Vertfcstion: Paxhr Open
Push
?
A
hittien
"unbokked"
Unlocked
Locked
Neglect locked unfocked ca86a in statistical
Couldin explore door locked case
Curius straiegy
Undertying
Causal Rules
Problem 1
Problem
Curious Causality-Sookingg Agents
Figure 1: Illustration of the Meta-Causa Graph concept and the Curious Sausality-Seekingg 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 Duusllty-Seekingg Agents actively intervene to verify causal relationships and discover the critical meta states that determine when causal structures change, enabling the agent to build comprehensive
Graph without requiring predefined domain labels.
ground-truth causal laws remain unchanged, the observational data is collected through local, context- limited exploration. Under conditions of partial observability, this invariance becomes obscured within a limited observational window, giving rise to the illusion of shifting causal relationships rule valid in one context may break down in another. For instance, the "push - open" relation holds for an unlocked door, but this relation vanishes when the door is locked. In this case, the same action produces effect, leaving the true causal mechanism ambiguous. Traditional causal modeling in such scenarios faces two main challenges: (1) methods that learn a single uniform (time series- -rries-based) causal graph neglect these context-dependent changes [47), and (2) lack of active curiosity-drive exploration strategy, which makes it struggle to adapt to discover the global causal rules in rich, unseen variations of potential open-ended worlds 29 [T].
To deal with these two challenges, we propose the Meta-Causal Graph to describe the world, which unified structure that captures how causal relationships evolve across different states. Rather than relying on a uniform graph, the Meta-Causal Graph is a minimal representation that contains particular
multiple causal subgraphs, each corresponding meta states "trigger" the activation of the appropriate subgraph, allowing us to model shifts in causal influence, such as when locked door severs the "push - open" link. We introduce framework termed Curious Causality-Seeele Meta-Causal Graph. The main contributions of this work are given below:
state. Transitions between these
Agent, which actively explores interventions to uncover the
Theoretically, we derive sufficient conditions ensuring that strategically chosen families of interven- tions uniquely identify meta-causal graphs, overcoming observational limitations.
Methodologically. we design a framework which leverages a curiosity -drivee exploration strategy where agents selectively intervene in the environment to learn causal world model. First, employ interventional verification to directly test whether specific variables causally influence others under different state conditions. Unlike methods relying solely on observational data, our approach actively intervenes on variables to establish causal links. Second, we introduce a targeted exploration strategy designed to discover meta states, critical configurations where causal structure changes.
Empirically, our approach demonstrates substantial improvements over purely observational baselines in simulated benchmarks, showcasing its capability to accurately recover complex causal structures. By explicitly modeling and reasoning about these causal relationships through active interventions, our proposed method significantly enhances the robustness and adaptability of world models.
2 Preliminaries
In this work, we employ causal modeling to construct world models. We begin by introducing fundamental definitions and terminology related to causal graphs and interventions.
Causal Graph. A causal subgraph is formalized as a directed acyelic graph (DAG) g (pp) F.), where the vertex set [p] {1,
................................ 2..,....,
indexes collection of random variables X (X1) with

---

joint density f(X). The edge set E' S [p] x [p] encodes direct causal relationships between variables. For each node E [p), we write (aş for its set of parents. An edge + E indicates that X; exerts direct causal influence on Xj.
Intervention. Directly learning the causal structure from observational data is often insufficient, as it can lead to multiple possible causal graphs that share the same v-structures. To resolve this ambiguity, we can use interventions to break the observational ambiguity. Intervention is powerful tool in causality. An intervention refers to the manipulation of variable within a causal model, often represented by the do-operator. For example, consider intervening on state Xg. where the system originally takes the value Xx Tk. An intervention is performed by setting the state variable to new value, as in do(X; = z²). such as setting "door open" to true in the environment.
During world-model learning, we perform sequence of environment interventions to reveal causal structure. To better formalize the intervention influence, this paper, we introduce the notion of an intervention target. At intervention step k, we select single variable X; as the target, which we denote by the intervention target I = fi), where E [p]. Applying an intervention do( = x') removes all incoming causal links into X; Concretely, if g is the original DAG, the intervention graph ((x) is obtained by removing every edge of the form j + i. Vj E Pag(i), from G. This modified graph accurately reflects the altered generative process under the do-operation on X,,
Definition 1 (Intervention graph (27). Let D (Ip) F.) be DAG with vertex [p) and edge set E, and I C [p] intervention target. The intervention graph of D is the DAG D(T) ([p), E(D), where E()) {a + + e $ I}.
Given I, the Interventional Markov Equivalence is defined as follows:
Definition 2 (Interventional Markov Equivalence (27)). Two DAGs Dj and D2 are interventional Markov equivalent given set of intervention larget I if D1 and D2 have the same skeleton and the same v-structures, and (D) and D have the same skeleton and the same v-structures for all I.
Then the Interventional Markov Equivalence Class (I-MEC) is the set of DAGs that are interven- tional Markov equivalent to each other given a set of intervention targets I. Given a set of intervention targets I, we can determine the I-MEC of the system. However, the DAGs in the I-MEC are not unique. We denote the set DAGs in the I-MEC as Dz. Although DAGs Dz share the same skeleton and v-structures, they may contain edges with different directions.
3 Meta-Causal Graph: Definition and Identifiability
In this section, we present our Meta-Causal Graph for a world formulation. We begin by defining the Meta-Causal Graph, which is compact causal graph comprising multiple context-specific causal subgraphs, varying from latent meta states in Section 3.1 Then, we show the identifiability of meta states Section 3.2 and their indicated causal subgraphs in Section 3.3 by interventions.
3.1 Meta-Causal Graph
We consider an environment whose states are represented by a set of variables X (Xe E X, such that the "door open" is state in the environment. X in the environment can be intervened on by manipulating its values directly or influenced by taking actions indirectly. Then, the Meta-Causal Graph is defined on the environment state follows:
Definition 3 (Meta-Causal Graph). The Meta-Causal Graph MG consists of collection of causal subgraphs MG = { (Gu]weu where each Gu corresponds to a distinct meta state u E U.. The causality skeleton matrix M,, is binary matrix where ji a parent of variable j (i.e. E Pa(j)) in the causal graph Gw. A state-to-meta state mapping C: X U assigns each system observation x E X to its active meta state C(x). We call Mg = ((Gu),C) a Meta-Causal Graph if the following conditions are satisfied:
indicates that variable is
1. There exists a ground truth mapping C X + determining the real meta states.
2. Causality skeleton matrices of causal subgraphs are sufficiently different. For all a + w', M., # Mas, i.e., there exists index (i.j) such that M [i & Mv [i,
3

---

The underlying Leta-Causal Graph in the environment is assumed to be invariant; however, its complete identification requires two steps: (1) identification of latent meta states, and (2) identification of the causal subgraph corresponding to each meta state.
3.2 Identifiability of Meta States
In this subsection, we will discuss the identifiability of the meta states from the environment. We first introduce the concept of swap-label equivalence between two mappings.
Definition 4 (Swap-Label Equivalence). Two mappings C X + Uj and C2 X + U2 are swap- label equivalent if there exists permutation U1 + U2 such that C2(x) = ((C) Væ E
The definition indicates that if we can identify the meta states up an permutation transformation, we say that the meta states are identifiable up to label swapping. We then show that the meta states are identifiable up to swap-label equivalence under the following assumption.
Assumption (Mixed Data Structure Learning). Let MG = [Gi x be a Meta Meta-Causal Graph with K causal subgraphs corresponding to distinct meta states E U. Consider a dataset 'D where each sample is generated from one of these causal subgraphs. Let Sp S (1,2..., K} denote the indices of causal subgraphs that actually contributed samples D. When learning a single causal graph & from the mixed dataset D (treating all samples as if they were generated from a single causal subgraph), the estimated parent set for each variable X; in 9 satisfies:
(ag(X)) = U Pag. (X²) Vj [p).
sESp
Under this assumption, the unified graph from pooled data contains the union of all direct parent-child edges present in each causal subgraph. Consequently, the latent meta state are identifiable as follows.
Theorem (Identifiability of Meta States). Under Assumption 1 the learned mapping C: X - t swap-label equivalent to the ground truth mapping C: X -UU.
Identifiability when number of meta states is unknown: In practice, we often overparameterize the meta state space by choosing more clusters than actually exist. We prove that this overparameterizztion does not harm identifiability; even with excess clusters, each true causal subgraph can still be recovered. The following theorem shows that, provided the number of clusters is sufficiently large, the underlying subgraph structures are identifiable up to the observational equivalence.
Definition 5. Two mappings C₁ : - U, and C2 + U2 are observationally equivalent if there exists mupping &b U1 + U2 such that C2(x) = (CC1 (x)) for all x E X.
Intuitively, observational equivalence refers to the condition where two mappings from states to meta states produce the same partitioning of observations, possibly up to relabeling.
Theorem (Identifiability of Overparameterized Meta States). Under Assumption | if the learned mapping C : X + Ú satisfies |Ú] > |U', the learned mapping C' is observationally equivalent to the ground truth mapping C : X - U.
3.3
Identification of Causal Subgraphs
In this subsection, we demonstrate that causal subgraphs for each meta state become uniquely identifiable once edge directionality is determined through interventions.
3.3.1 Causal Subgraph Identification
We employ multiple interventions to determine the causal structure of environment, represented as intervention target indicator set I = { K where I C {2]} comprises the indices of intervened variables selected from state variable set X {Xikeen at intervention step k. These intervention sets are collected from historical data or agent interactions. We establish sufficient conditions ensuring causal graph edge identifiability through appropriate interventional selection.
Proposition (Identifiability of Causal Subgraph). The causal subgraph D can be uniquely identified if there exists an intervention target indicator set such that for all edges a -> E D, there exists an intervention target I E I such that I n {a, = 1.
4

---

This proposition establishes that causal subgraph identifiability requires edge-specific interventions targeting single connected variables, enabling directional determination.
Curious Causality-Seekinng Agents for Open-Ended World Modeling
We propose a novel framework to learn the Meta-Causal Graph through Causality-Seeking Agent. Our approach comprises three core components: (1) curiosity-drive interventional exploration in an open-ended world, (2) Meta-Causal Graph discovery from agent experience, and (3) continual world model learning and updating. Together, these components enable the agent to actively discover and refine causal structures in open-ended environments.
4.1 Curiosity-Drive Interventional Exploration in Open-Ended Worlds
Accurate identification of the Meta-Causal Graph cannot rely solely on passive experience. First, causal structures inferred from purely observational trajectories often suffer from edge misorientation. Second, an agent's s experience typically spans only limited subset of world states, leaving parts of the causal graph unobserved. To reveal the full Meta-Causal Graph in open-ended environment, the agent must complement passive observation with targeted, curiosity-driven interventions that actively probe uncertain causal relations.
Curiosity Reward. Our framework treats curiosity as a general intrinsic signal that can be instanti- ated in multiple forms. Specifically, experimented with the following formulations:
(I) Edge-Entropy. 1e Eijebl II (Mc(x,)) (i),)))) where Mc(x))1..1 posterior probability of edge from variable to in the current meta state C(X,), and H(-) is the Shannon entropy. This term prioritizes interventions on the most uncertain parts of the causal graph.
[0, 1] denotes the
(2) Prediction Uncertainty. Iime If [po(I++1 It, (t),, where pe(1++1 It, de) is the world model's 's predictive distribution over the next state given current itateacttion pair (Iz, oz). Higher entropy indicates epistemic uncertainty about future transitions.
(3) Feature Discrepancy. If E(*++1) - Es where E(-) is the learned feature encoder and Eg 1#(++1) is the predicted feature of the next state 4441- This term measures reconstruction error in the feature space, highlighting dynamics that the model fails capture.
(4) Predictive-Distribution Discrepancy. I log po negative log-likelihnn likelihoood of the observed transition.
It; 0t), which the surprisal or
Any of these intrinsic terms can be used the curiosity reward Rg select interventions, and in our main results we adopt the edge-entropy variant as the default while other forms yield comparable behaviors (see Table
Intervention Verification. After executing curiosity-guic -guided interventions, we validate and refine the learned causal structures by directly estimating causal effects. For each variable X,, we perform interventions do(),=x)) and observe the resulting changes in other variables. We estimate the causal effect of X, on X₁ as
D = log P(X |X do(X,=x))) - log P(X++1 |xt).
We integrate these effects into mask-refinement loss
Lmask ーλ
Σ
log Mij + X2
Σ
log Maj.
((..):)A.) >
where 7 2 0 controls sensitivity.
Interventional Reachability. In realistic environments, not all variables are directly intervenable. To formalize which causal relations are identifiable under such constraints, we introduce interventional reachability. We define an intervention operator F E {0. NxN where F(k, = 1 if and only state can be obtained from state by a single allowed intervention, and transition operator T {0, 1J for the environment's natural dynamics after intervention. The composite TF
5

---

models one cntveentioo-transsite cycle, and nonzero entries of (TF)* enumerate states reachable
within cycles. state reachable if and only if some finite h satisfies [(TF)* (7))=z)) 20
[
0. These
constraints define the feasible set of causal interventions. Details are given Appendix B
4.2 Meta-Cause Graph Discovery from Agent Experience
Given trajectories collected through curiosity-driven interventions, the agent infers structured representation of causal dependencies that vary across latent meta states. The goal is to jointly learn (i) the discrete latent variable C(2)) indicating the current meta state, and (ii) the corresponding causal skeleton matrix Mc(x,) that governs the local dynamics.
Representation Learning. To model context -eeeeden causal dependencies from trajectories generated by curiosity-driven interventions, we discretize the latent dynamics of the environment using a cuttor-qaatized variational auto-encoder (VQ- (Q-V [I. The encoder maps concatenated state-action pairs (Ix, ax) into latent vector Re, which is quantized to the nearest codebook entry &u. associated with a meta state u. Each codebook entry defines a distribution over causal structures, from which the causal skeleton matrix Mc(x.) is sampled. Further details are provided in Appendix
Effect of Lossy Representation. The latent representation learned through vector quantization is inherently lossy due to imperfect mapping. Such lossy representations may distort the underlying causal factors and thus influence the identification of both meta states and their associated causal subgraphs. We provide an error bound of lossy representation in the following.
Proposition (Effect of Representation Accuracy on Misclassification). Given a set of ground-truth meta states = {u] and corresponding codebook entry index Ú = { let pk = Emeo PuPE denote the probubility thut sample x is mapped to codebook entry zág- Then, under the lossy representation induced by encoder, the probability of misclassificution
P(misclassification)
= =1 -1--Em
Mup%(1 (1- - Pk + ph.
WEU
When n(ph + PhH)) < 1, this can be approximated as
P(misclassification)
Σ
Pk (HuPtt)
Adaptive Codebook Fusion. To prevent redundant meta states resulted from overparamterization and maintain a minimal representation, we adopt an Adaptive Codebook Fusion mechanism. During training, codebook entries encoding similar latent embeddings and producing comparable decoded causal skeleton matrices are automatically merged. This fusion step is integrated into the quantization update, requires extra supervision, and reallocates capacity for novel meta states.
4.3 Transition Probability Learning
The previous subsections describe how identify the parent set of each state variable, Pag, (X,). using the learned causal skeleton matrix Mg. In this section, we demonstrate how to leverage the confirmed causal skeleton to quantify inter-variable causal effects, thereby grounding and validating the transition probabilities of the world model for next-state prediction. We train our model by maximizing the log-likelihood of the observed sequences, thereby learning the transition dynamics.
LMeE = log P(X |Pag.. (X₃)).
The complete world model learning objective. To encourage the discovery of parsimonious causal structures that capture essential relationships while avoiding spurious connections, we incorporate a sparsity regularization term: Lxparse = || II, where || denotes the LI norm of the causal skele- probability matrix, promoting sparsity in the learned causal graphs. The complete optimization objective integrates multiple components is as follows:
L = LMLE + Aquarse Eyparse + Amask Lmask + Aquamtizaation Segaattttiinnuuuuatttiiiiiiooon
Training proceeds by alternating among interventional exploration, causal subgraph updating, world model learning, and codebook refinement. Algorithm summarizes the complete learning procedure.

---

Table 1: Prediction accuracy on OOD states in the Chemical environment. The number of noisy
nodes in the downstream tasks is denoted as n. All values are reported in percentage (%)
Algorithm
Fork
Chain
D=2
D=
D=0
n==2
n=
D=
GNN
36.29:1:345
25.80=33448
21.58:13.44
29.22:13.39
23.28:4.98
20.53:16.96
MLP
31.1111699
30.44=2.25
32.39.1..76
28.66=3.65
26.5214.26
24.15:4.17
NCD
41.60455008
37.47=2.13
42.2741182
40.04=6.21
37.4742.98
41.19:1..66
FCDL
57.82-49.00
49.29=8.90
47.70:16.68
10.66+10.10
48.81
48.05:55.88
Modular
26.534-3.45
24.73-4.661
26.7348.31
25.2444.68
24.944-4.81
25.09+55.91
NPS
40.56:24.61
26.8144.37
23.0244.27
38.73=2.63
27.6944.228
24.4543.84
CDI
35.59+1.85
42.2241..39
34.90+1.59
36.5241.72
4..06++1.29
GRADER
37.93+1.066
38 1.694
45.74:22.25
36,82=3.122
37.41+2.84
43,48+4.11
Oracle
33.87++1.34
1.800
42.47:+0.75
34.63-11.3 1.78
38.31++2.48
42.87-+2.08
Sandy-Mixure
31.
32.4710.0161
33.7244.11
30.08=3.12
29.3115.18
27. 27.43:22.20
Transformer
25.13:1:0.33
24.37:22.58
21.90:12.35
29.62::0.65
30.37:12.81
29.78
MCG (wo mask loss)
58. 18.15
51.70=8.58
46.04.18.56
47.75 .2
46. 16.8977.87
47, 47,94:16.60
MCG (wo intervention verification)
48.28-18.15
5.04
46.54443.72
8..36:8805
48.72:+-4.42 6.83
MCG (ours)
63.18:113.94
50.47=9.87
50,04:8.56
61.99:6.58
49.782..11
49.69-15.14
Table 2: Average episode rewards on downstream tasks for each environment. The number of noisy nodes introduced in the downstream tasks is denoted
Chain
Fork
Algorithm
Magnetic
n=2
n=4
n=6
n=2
n=4
n=6
GNN
6.89:10.22 28
6.38:10.88
6.56
6.61 0.992
6.15
6,95
2.23 0.90
MLP
7.39++0.65
6.63=+0.58
6.78+0.93
6.49++0.88
5.93=0.71
6.84+1.17
2.10
NCD
9.60+1.52
8.86++0.23
10.32+0.37
10.95=1.63
9.11+0.63
9.11+0.63
+
FCDL
11.16:1.5
10.39:22.84
10.62:42.22
13.98=2.01
13.36:2.14
12.91+2.40
0
Modular
6.61#:0.63
7.01=0.55
7.04=1107
6.05=
5.65=0.50
6.43=1.00
0.88 =
NPS
6.92.11.03
6.88.10.79
6.8010039
5.82.10.83
5.7510.57
5,5410.80
0.91 1.669
CDL
8.71+00555
8.65
10.23:10.50
9.37+11.3
8.23+0.40
9.50
1.10
Oracle
8.47:10.69
8.85
10.29
7.83110.77
8.04
9.66
0.95 0.55
Sandy-Mixture
6.81.10.17
6.73.00.20
1.02
6.95110.2
6.7110.20
7.03
1.63 0.02
Transformer
6.45++0.28
6.73+00118
7.31++0.33
6.54+00118
6.68+0.27
7.00+0033
2.13 0.01
MCG(ours)
13.82-33884
12.49+2.39
12.45+1.37
14.65-2.75
14.06=22.64
13.28+22.04
3.19 0.14
5
Related Work
World Models. World models (21
62 70 enable agents to summarize past interactions, predict
future states in the environment, and evaluate candidate actions [23 21
55
26 25). However,
methods that rely solely statistical correlations often break down when environmental conditions
shift, undermining their ability to generalize robustly 55
18 19
Causal Discovery for Open-Ended World. Causal discovery provides rigorous framework for modeling the generative processes that govern complex systems, revealing how one variable brings about changes in another. By explicitly representing cause-effect relationships, causal methods yield more compact and invariant descriptions of reality than do purely statistical correlations
36 44 structure, which aims to recover causally meaningful latent factors and structures, enabling invariance and transfer in downstream tasks [64 representation learning focusses on learning a static causal structures passively from observational
81
8 9 15 45). Causal representation learning provides useful tool to learn causal
39 80 79
78 However, most of the work of causal
data.
However, identifying causal structure from observational data alone is a challenge since multiple directed acyclic graphs can entail the same conditional independencies and v-structures (76). To
resolve these ambiguities, researchers have turned to the collection of active causal interventional
data by actively perturbing variables distinguish among candidate graphs [52 33 72 6 66). More
recent work has even addressed settings with unknown intervention targets 27 28
Despite these advances, most causal discovery algorithms presume a single uniform, context- 7 in the world. Some recent methods have 34 but they typically
independent causal graph 47 sought to learn causal structures in changing world [57]
82
10 75
58
49
32
31
7

---

Acrunecy
Copo Sizz
Code Suz
Code
Code Nio
Oode 3100
(a) Prediction accuracy on Chain.
(b) Prediction accuracy on Fork.
n-44
n-4
Epiionde Towwent Moan
15
Expaande Rewverd
Numberol Mata statea
Namber Staten
Number
Mara Siste
Mumber
(c) Downstream rewards Chain.
(d) Downstream rewards on Fork.
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.
(a) Pattem (h) Pattern
Pattern 1 (d) Pattem 2 (e) Pattern 3 (f) Pattern 4 (g) Freq. (h)
Freq.
(k=2)
(k=2)
(k=4)
(k=4)
(k=4)
(k=4)
dist. (k=2) dist. (k=4)
Figure 3: Comparison of causal patterns discovered with different numbers of meta states. When the number meta states is set to 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 patterss share similarities (Patterns and 3), while others occur with lower frequency (Pattern 4) as shown in the frequency distributions (g.h).
require strong prior knowledge [57 29 complete access to known dynamic model [10] and they lack active, curiosity-driven exploration strategies for open-ended environments [56].
We introduce intervention-drivee framework that empowers agent actively probe its surround- ings order to uncover context-dependent causal graphs. By combining targeted interventions with a curiosity-dii driven exploration policy, our approach adaptively reveals the global causal rules governing rich, previously unseen environments. Extended related work could be found Appendix C
6
Experiments
We evaluate our method by examining the following research questions: (1) Does our method learn more accurate causal world model? (Table (2) Does it enhance performance on downstream tasks? (Table 2) (3) How does overparamettriizatio affect our method's performance? (Figure (4) Do the curiosity-driven reward and intervention verification improve the learned causal world model quality? (Table Table 2) (5) Does our method enable more accurate causal subgraph learning? (Figure
6.1 Experiment Setup
We compare our method against several baselines: MLP, a dense model that predicts transition dynamies p( (ate; at); Modular, a modular network with separate modules for each state vari- able; GNN [4)] graph neural network for predicting transition dynamics; NPS (2), which learns sparse and modular dynamics; CDL [77), which learns a static causal model from data;
GRADER [13,, which learns static causal model using conditional independence tests; Sandy- Mixure (5). which uses Jacobian matrices of MLP layers to learn the local causal graph: Trans- former which infers causal graphs by attention: NCD [31] neural causal discovery algorithm learning causal graphs for each sample; and FCDL [3 learns causal graph facilitate the
robustness for reinforcement learning. Please check Appendix B for details of the experimental setup.

---

Algorithm Curious Causality-Seeking Meta Causal World Modeling
Require:
Observational states X, meta state embedding space U, initial codebook embeddings {ew
encoder E, decoder D. hyperparameters Xqpars; Amusk, Aquamizzation
1: Initialize encoder parameters Φ, decoder parameters 0, and embeddings { (zuhueu
2: while not converged do
3: Compute embedding assignment:
arg yeu (X)x(X) zul12
4: Compute causal skeleton probabilities matrix: My + D(e.)
5: Sample causal skeleton matrix MM. j]] GumbelSoftmax [M 11,
6: Select intervention do(X; z maximizing reward R& in Section 4.1 Perform intervention and record resulting state transitions
7:
8:
Estimate causal effects Ayy from interventions and update causal mask parameters via Lmask
9: Predict next state x from parent set Pag. (X,) and update world model via LMLL
10: Encourage sparsity in causal graph via Esparse
11: Update embedding codebook {E.] via Lquamtization
12:
end while
13:
return Learned causal subgraphs Iu, world model parameters, and embeddings {eu)
6.2 Environments
Causal Graphs
Chain
Chemical 39 We use the Chemical environment to evaluate the performance of the proposed method on learning the causal graphs in system with multiple causal structures. There are several causal graphs (full, fork, chain) in the Chemical environment and the causal graph depends on the state of the objects. We use two set- tings of the Chemical environment: and (2) full-chain (Chain). Magnetic (31) The Magnetic environment is a robot arm manipulation task (Figure 4 This environment is built based on the Robosuite suite. Please refer to Appendix B for more details.
C-O-O-O
Fork
Full
Fork
(Fork)
6.3 Experiment Results
Prediction Accuracy. We evaluate the prediction ac- curacy of our proposed method in comparison to several baseline approaches. To systematically assess robust- ness, we introduce varying levels of noise to the state information. For the full-fork and full-chain tasks, we corrupt the values of 2, 4, and 6 nodes, respec- tively, to measure performance degradation under increasing noise. In the magnetic environment, we assess prediction accuracy when box position coordinates and ball/box color properties are corrupted, thereby testing the model's performance under partial observability. All experiments are conducted eight times, with results averaged. As shown in Table I our method consistently achieves the highest prediction accuracy across nearly all baselines, demonstrating its effectiveness in learning a more accurate causal world model.
Figure 4: Visualization of environments: (top) Chemical; (bottom) Magnetic.
Task We evaluate our method against basclines on the Chemical downstream task, using the trained model to predict future states with the learned model for decision making. As shown in Table 2 our approach achieves the highest reward, demonstrating its effectiveness in learning a more accurate causal world model.
Downstream
Overparameterrzation We evaluate our method using 2, 4, 8, and meta states, with results presented in Figure 3 Performance declines markedly when using fewer meta states than actually exist, as the model fails to capture distinet causal structures. Conversely, overparameterized models consistently outperform underparameterized ones, empirically validating Theorem 2 The learned causal graphs with codebook sizes of 2 and 4 reveal that with 2 meta states, the model learns two distinct causal graphs (Pattern and Pattern 2) corresponding to the two meta states, consistent

---

(a) Causal graph of Full.
(b) Causal graph of Chain.
Figure Comparison of learned causal graphs in the Chemical task: left left-MG (ours); mid- ree-reference (FCDL); right- -ground truth.
with Theorem I With 4 meta states, some codes correspond to similar underlying causal graphs (as shown by the similarity between Pattern and Pattern 3), while codes representing significantly different causal graphs (such as Pattern 4) occur less frequently. minimally affecting overall model performance. This demonstrates the model's capacity consolidate redundant representations when overparameterized, aligning with our theoretical findings.
Ablation Study We conduct ablation studies to evaluate the effectiveness of the proposed method. We remove the following components from the proposed method: (1) the causal loss Luusk and (2) the reward function. The first ablation study aims test the effectiveness of the intervention verification, while the second ablation study aims to test the effectiveness of the active intervention exploration. The results are shown in Table 2 Overall, the proposed method outperforms the ablation methods. The performance drop of removing the causal loss Lmikk indicates that the incorporation of the causal loss can help the model learn a more accurate causal world model. The performance drop of removing the reward function indicates that the auriosity-based reward function can help the agent learn more accurate causal world model through active exploration.
Partial Observation We modify the Chemical environment to explicitly simulate challenges arising from limited observation windows. This experimental setup evaluates how method performs as new variables gradually become observable. The results indicate that our approach remains robust under partial observability and shows strong potential for handling more realistic, dynamically evolving environments. The results are shown in Table g
Efficiency We compare baselines with longer training rounds and more parameters to show efficiency of our method. Results are reported in Table 10 and Table IT
Causal Subgraph We visualize the learned causal subgraphs in Figure 5 It is evident that the proposed method successfully captures the underlying causal structures of the environment.
Additional results are in Appendix B and the source code is included in the supplementary materials.
7 Conclusion, Limitation and Future Work
In this work, we introduced a method model the environment using the Meta-Causal Graph that explicitly captures how causal relationships evolve across different environmental montexts-meta state. Our approach addresses two critical limitations in existing causal world models: the inability of uniform causal graphs to capture context-depende hard to learn the open-world environment. We theoretically established the identifiability of meta states and their corresponding causal subgraphs, and developed the Curious Cauuality-Seekiii framework that actively explores environments through interventions guided by a curiosity-driven reward function. Our empirical evaluations demonstrated that our method outperforms existing approaches across both synthetic tasks and a challenging robot arm manipulation task.
changes and the lack of active exploration making
Agent
Limitation and Future Work Despite these advances, real-world exploration remains subject to inherent constraints: certain states are fundamentally unreachable via any sequence of interventions, and the implications of such unobservability for causal inference have yet to be analyzed (see Appendix A.8). Moreover, interventions incur costs, and practical agents must operate within limited budgets. Future research will investigate the impact of unreachable states on the identifiability and accuracy of learned causal structures, and will develop budget-aware exploration strategies to optimally allocate limited intervention resources.
10

---

References
[1] Pranav Agarwal, Sheldon Andrews, and Samira Ebrahimi Kahou, Learning to play atari in world of tokens. arXiv preprint arXiv:2406.01361, 2024.
[2] Anirudh Goyal ALIAS PARTH GOYAL, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael C Mozer, and Yoshua Bengio. Neural production systems. Advances in Neural Information Processing Systems, 34:25673-25633 2021.
[3] Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos J Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. Advances in Neural Information Processing Systems, 37:58757-58791 2024.
[4] Amir Bar, Gaoyue Zhou, Danny Tran, Trevor Darrell, and Yann LeCun, Navigation world models. arXiv preprint arXiv:2412.03577. 2024.
[5] Stephan Bongers, Tineke Blom, and Joris M Mooij. Causal modeling of dynamical systems. arXiv preprint 18X3::880330874 2018.
[6] Philippe Brouillard. Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems, 33:21865-21877, 2020.
[7] Hongye Cao, Fan Feng, Meng Fang, Shaokang Dong. Tianpei Yang, Jing Huo, and Yang Gao. Towards empowerment gain through causal structure learning in model-based reinforcement learning. In The Thirteenth International Conference on Learning Representationss 2025.
[8] Weilin Chen, Ruichu Cai, Zeqin Yang, Jie Qiao, Yuguang Yan, Zijian Li, and Zhifeng Hao. Doubly robust causal effect estimation under networked interference via targeted learning. In International Conference on Machine Learning, pages 6457-6485. PMLR, 2024.
[9] Fengxiang Cheng, Haoxuan Li, Fenrong Liu, Robert Van Rooij, Kun Zhang, and Zhouchen Lin. Empowering iims with logical reasoning: A comprehensive survey. In Proceedings of the Fhirth-Fourth International Joini Conference on Artificial Intelligence, 2025.
[10] Rohan Chitnis, Tom Silver, Beomjoon Kim, Leslie Kaelbling, and Tomas Lozano-Perez. Camps: Learning context-specific abstractions for efficient planning in factored mdps. In Conference on robot learning, pages 64-79, PMLR, 2021.
[11] Etched Decart, Q McIntyre, S Campbell, Xinlei Chen, and Wachen. Oasis: A universe in a transformer. URL: patps://oasismmode github. io, 2024.
[12] Antoine Dedieu, Joseph Ortiz, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, Swaroop Guntupalli, Miguel Lazaro- -reaill,,,,,,,a, and Kevin Patrick Murphy. Improving transformer world models for data- cntt-eefficentt rl. arXiv preprint XrXiv:2502.01591, 2025.
[13] Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Generalizing goal-conditioned reinforce- ment learning with variational causal reasoning. Advances in Neural Information Processing Systems, 25:26532-26545 2022.
[14] Gonçalo Rui Alves Faria, Andre Martins, and Mário AT Figueiredo. Differentiable causal discovery under latent interventions. In Conference Causal Learning and Reasoning, pages 253-274, PMLR, 2022.
[15] Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for non- stationary reinforcement learning. S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 31957-31971. Curran Associates, Inc., 2022.
[16] Ruili Feng, Han Zhang, Zhantao Yang. Jie Xiao, Zhilei Shu, Zhiheng Liu, Andy Zheng, Yukun Huang, Yu Liu, and Hongyang Zhang The matrix: anfnitt-horizon world generation with real-time moving control. arXiv preprint arXiv:2412.03568, 2024.
11

---

[17] Xidong Feng, Yicheng Luo. Ziyan Wang, Hongrui Tang, Mengyue Yang, Kun Shao, David Mguni, Yali Du, and Jun Wang. Chessgpt: Bridging policy learning and language modeling. Advances in Neural Information Processing Systems, 36:7216-7262, 2023.
[18] Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Charles Blundell, Sergey Levine, Yoshua Bengio, and Michael Curtis Mozer. Factorizing declarative and procedural knowledge in structured, dynamical environments. In International Conference on Learning Representations, 2021.
[19] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sengey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1900) 10893, 2019.
[20] Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research, 19:399-468 2003.
[21] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 2018.
[22] Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprini 1rXiv 2019.
[23] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02199. 2020.
[24] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint 200V::2300.040044 2023.
[25] Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for continuous control. arXiv preprint 100. 2023.
[26] Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint WrXiv:2203.04953 2022.
[27] Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional markov equivalence classes of directed acyelic graphs. The Journal of Machine Learning Research,
22644 2012.
[28] Alain Hauser and Peter Bühlmann. Jointly interventional and observational data: estimation of interventional markov equivalence classes of directed acyclic graphs. Journal of the Royal Statistical Society Series B: Statistical Methodology, 77(1):291-318. 2015.
[29] Biwei Huang, Kun Zhang, Jiji Zhang, Joseph Ramsey, Ruben Sanchez Romero, Clark Glymour, and Bernhard Schölkopf. Causal discovery from heteengenneous/nonsstaioonn data, Journal of Machine Learning Research, 21()) 1-33, 2020.
[30] Edward Hughes, Michael Dennis, Jack Parker- Holder, Feryal Behbahani, Aditi Mavalankar, Yuge Shi, Tom Schaul, and Tim Rocktaschel. Open-endedness is essential for artificial superhu- man intelligence. arXiv preprint arXiv:2406.044 042600 2024.
[31] Inwoo Hwang, Yunhyeok Kwak, Suhyung Choi, Byoung- Tak Zhang, and Sanghack Lee. Fine- grained causal dynamics learning with quantization for improving robustness in reinforcement learning. arXiv preprint 2tt.:2440600222344 2024.
[32] Inwoo Hwang, Yunhyeok Kwak, Yeon-Ji Song, Byoung-Take Zhang, and Sanghack Lee. On discovery of local independence over continuous variables via neural contextual decomposition. In Conference on Causal Learning and Reasoning, pages 448-472. PMLR, 2023.
[33] Antti Hyttinen, Frederick Eberhardt, and Matti Järvisalo. Constraint- based causal discovery: Conflict resolution with answer set programming. In UAI, pages 340-349, 2014.
[34] Fateme Jamshidi, Sina Akbari, and Negar Kiyavash. Causal imitability under context-specific independence relations. Advances Neural Information Processing Systems, 36:268100268833 2023.
12

---

[35] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with softmax. arXiv preprint arXiv: 1611.01144 144, 2016.
[36] Dominik Janzing and Bernhard Schölkopf. Causal inference using the algorithmic markov condition. IEEE Transactions on Information Theory, 36(10):516-5511 2010.
[37] Jiarui Jin, Yuwei Wu, Haoxuan Li, Xiaoting He, Weinan Zhang, Yiming Yang, Yong Yu, Jun Wang, and Mengyue Yang. Large language models are demonstration pre-selectors for themselves. In Forty-second International Conference Machine Learning. 2025.
(38] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model- based reinforcement learning for atari. arXiv preprint arXiv: 00374, 2019.
[39] Nan Rosemary Ke, Aniket Didolkar, Sarthak Mittal, Anirudh Goyal, Guillaume Lajoie, Stefan Bauer, Danilo Rezende, Yoshua Bengio, Michael Mozer, and Christopher Pal. Systematic evaluation of causal discovery in visual model based reinforcement learning. arXiv preprint arXiv:2107.00845 2021.
[40] Thomas Kipf, Elise Van der Pol, and Max Welling, Contrastive learning of structured world models. arXiv preprint arXiv:1911.122444 2019.
[41] Abhinav Kumar and Gaurav Sinha. Disentangling mixtures of unknown causal interventions. In Uncertainty in Artificial Intelligence, pages 2093-2102. PMLR, 2021.
[42] Tobias Lang, Marc Toussaint, and Kristian Kersting, Exploration in relational domains for model-based reinforcement learning. The Journal of Machine Learning Research, 13(1):3725- 3768, 2012.
[43] Anson Lei, Bernhard Schölkopf, and Ingmar Posner. Causal discovery for modular world models. In NewrlPS 2022 Workshop on Neuro Causal and Symbolic Al (nCSI), 2022.
[44] Yunzhu Li, Antonio Torralba, Anima Anandkumar, Dieter Fox, and Animesh Garg. Causal discovery in physical systems from videos. Advances in Neural Information Processing Systems, 33:9180-9192, 2020.
[45] Haohong Lin, Wenhao Ding, Jian Chen, Laixi Shi, Jiacheng Zhu, Bo Li, and Ding Zhao. Because: Bilinear causal representation for generalizable offline model-based reinforcement learning. arXiv preprint arXiv:2407.1096) 2024.
[46] Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. BISCUIT: Causal representation learning from binary interactions. In Robin Evans and Ilya Shpitser, editors, Proceedings of the Thirly-Ninth Conference on Uncertainty in Artificial Intelligence, volume 216 of Proceedings of Machine Learning Research, pages 1263-1273 PMLR, 31 Jul-04 Aug 2023.
[47] Yuren Liu, Biwei Huang, Zhengmao Zhu, Honglong Tian, Mingming Gong, Yang Yu, and Kun Zhang. Learning world models with identifiable factorization. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages
Curran Associates, Inc., 2023.
[48] Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In International conference on machine learning, pages 6348-6359. PMLR, 2020.
[49] Sindy Löwe, David Madras, Richard Zemel, and Max Welling, Amortized causal discovery: Learning to infer causal graphs from time-series data. In Conference on Causal Learning and Reasoning, pages 509-525, PMLR, 2022.
[50] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv: 1,00 2016.
[51] Christopher Meek. Graphical Models: Selecting causal and statistical models. PhD thesis, Camegic Mellon University, 1997.
13

---

[52] Nicolai Meinshausen, Alain Hauser, Joris M Mooij, Jonas Peters, Philip Versteeg, and Peter Bühlmann. Methods for causal inference from gene perturbation experiments and validation. Proceedings of the National Academy of Sciences, 1 113(77:7761-7368 2016.
[53] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models, arXiv preprint arXiv:2209 000088, 2022.
[54] Vincent Micheli, Eloi Alonso, and François Fleuret. Efficient world models with context-aware tokenization. arXiv preprint arXiv:2406. /9320, 2024.
[55] Tim Pearce, Tabish Rashid, Dave Bignell, Raluca Georgescu, Sam Devlin, and Katja Hofmann. Scaling laws for pre-training agents and world models. arXiv preprint arXiv:241 1.04434, 2024.
[56] Judea Pearl. Causality. Cambridge university press, 2009.
[57] Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems, 33:3976-3999) 2020.
[58] Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. Mocoda: Model-based coun- terfactual data augmentation. Advances in Neural Information Processing Systems, 35:18143- 18156, 2022.
[59] Jonathan Richens and Tom Everitt. Robust agents learn causal world models. In The Twelfth International Conference on Learning Representations, 2024.
[60] Stéphane Ross and Joelle Pineau. Model-based bayesian reinforcement learning in large structured domains. In Uncertainty in artificial intelligence: proceedings of the... conference. Conference on Uncertainty in Artificial Intelligence, volume 2008, page 476, 2008
[61] Reuven Rubinstein and Dirk P Kroese. The cross-entropy method: a unified approach to combinatorial optimization, Monte-Carke simulation and machine learning. Springer Science Business Media, 2004.
[62] Jürgen Schmidhuber. Making the world differentiable: on using self supervised fully recurrent neural networks for dynamic reinforcement learning and planning in non-stationary environ- ments, volume 126. Inst. für Informatik, 1990.
[63] Jürgen Schmidhuber. Gödel machines: Fully self-referential optimal universal self-improvers. In Artificial general intelligence, pages 199-226 Springer, 2007.
[64] Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 10965)x612-634, 2021.
[65] Julian Schritwieserr, loannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):644-60) 2020.
[66] Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal influence detection for improving efficiency in reinforcement learning. Advances Neural Information Processing Systems, 34:22905-22222 229188 2021.
[67] Sumedh A Sontakke, Arash Mehrjou, Laurent Itti, and Bernhard Schölkopf. Causal curios- ity: RI agents discovering self- self-supervised experiments for causal representation learning. In International conference machine learning, pages 9848-9858. PMLR, 2021.
[68] Peter Spirtes, Clark N Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2000.
[69] Chandler Squires, Yuhao Wang, and Caroline Uhler. Permutation- causal structure learning with unknown intervention targets. In Conference on Uncertainty in Artificial Intelligence, pages 1039-1048. PMLR, 2020.
[70] Richard Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4): 1991.
14

---

[71] Manan Tomar, Amy Zhang, Roberto Calandra, Matthew E Taylor, and Joelle Pineau. Model-invariant state abstractions for model-based reinforcement learning. arXiv preprint arXiv:2102.098550 2021.
[72] Sofia Triantafillou and Ioannis Tsamardinos. Constraint-based causal discovery from multiple interventions over overlapping variable sets. The Journal of Machine Learning Research, 16(1):2147-2205, 2015.
[73] Núria Armengol Urpí, Marco Bagatella, Marin Vlastelica, and Georg Martius. Causal action influence aware counterfactual data augmentation. arXiv preprint 2xXiv:445.188111 2024.
[74] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
[75] Zizhao Wang, Jiaheng Hu, Peter Stone, and Roberto Martín-Martin. Elden: Exploration via local dependencies. Advances in Neural Information Processing Systems, 36: 36:15556-15474 2023.
[76] Zizhao Wang, Caroline Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Building minimal and reusable causal state abstractions for reinforcement learning. In Proceedings the AAAI Conference on Artificial Intelligence, volume 38, pages 1777--17886, 2024.
[77] Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task-independent state abstraction. arXiv preprint arXiv:2206. 13452, 2022.
[78] Mengyue Yang. Quanyu Dai, Zhenhua Dong, Xu Chen, Xiuqiang He, and Jun Wang. Top-n recommendation with counterfactual user preference simulation. CIKM '21, page 2342-24-22 2352-2351,
2021.
[79] Mengyue Yang, Zhen Fang, Yonggang Zhang, Yali Du, Furui Liu, Jean-Francois Ton, Jianhong Wang, and Jun Wang. Invariant learning via probability of sufficient and necessary causes. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, pages 79832-79837 2023.
[80] Mengyue Yang, Furui Liu, Zhitang Chen, Xinwei Shen, Jianye Hao, and Jun Wang. Causalvae: Disentangled representation learning via neural structural causal models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, June 2021.
[81] Weiran Yao, Guangyi Chen, and Kun Zhang. Learning latent causal dynamics. arXiv preprint 2tti:22002008225 2022.
[82] Zhongwei Yu, Jingqing Ruan, and Dengpeng Xing. Learning causal dynamics models in object-oriented environments. arXiv preprint arXiv:2405. /2615, 2024.
[83] Amy Zhang, Zachary C Lipton, Luis Pineda, Kamyar Azizzadenesheli,, Anima Anandkumar, Laurent Itti, Joelle Pineau, and Tommaso Furlanello. Learning causal state representations of partially observable environments. arXiv preprint arXiv: 190...040.147 2019.
15

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims presented in the abstract and introduction accurately reflect the scope and contributions our work, which are supported by Section aadd Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer:
[Yes]
Justification:
We explicitly discuss limitations in Section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
16

---

Answer: [Yes]
Justification: We provide complete proofs for all theoretical results in Section 3 For Theorem [ and Theorem state the key assumption (Assumption and provide detailed proofs in Appendix A Similarly, for Proposition on subgraph identifiability, we list all assumptions and include complete proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper includes all necessary information for reproduction, including the algorithm pipeline [ the experimental setup, and the hyperparameters used in the experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibilitty In the case of closed-source models, it may be that access to the model is limited in
17

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We use publicly open source environments and release code in the supplemen- tary file.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
jublic/guides (aelli/cuuiies//ddeebbbaasssssss onlliicy)) for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips. //nips.c/ppie hwwpll/c/ggiiee/ooooeooooooo
ssi ssionpolicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We provide all the necessary details in Section and Appendix including hyperparameters and the type of optimizer used.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report the mean and standard deviation of prediction accuracy and standard deviation, and we provide the necessary information to understand how they were computed.
Guidelines:
The answer NA means that the paper does not include experiments.
18

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eπror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide the necessary information about the compute resources used in experiments in Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp:://oorrrs..c//u///////aaaa
Answer: [Yes]
Justification: We have reviewed and fully adhered to the NeurlPS Code of Ethics throughout the research process.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
19

---

Justification:
This work is primarily methodological and does not have direct societal
impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: This work does not involve the release of data or models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification; All assets used in this work, including datasets, code, and models, are publicly available. We have properly credited the original creators and ensured that all licenses and terms of use are respected.
Guidelines:
The answer NA means that the paper does not existing assets.
20

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaeed, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We release the code used in this work, and we provide detailed documentation on how use them.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
This work does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
21

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question:
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs used in this paper solely for text polishing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our L.LM policy for what should or should not be described.
22

---

A Theoretical Analysis and Discussion
This section establishes the theoretical foundations for identifying meta states and their corresponding causal subgraphs within our proposed framework. We present a comprehensive analysis of identifia- bility conditions under two distinct scenarios: swap-label equivalence and observational equivalence. The former addresses cases where meta states can be permuted without altering the underlying causal structure, while the latter cxamines situations involving overparameterization where the learned model contains more meta states than the ground truth. Additionally, we introduce intervention reachability to formalize which causal relationships identifiable via feasible interventions.
A.1 Summary of Key Results
Our theoretical analysis yields key results the identifiability of meta states and causal structures:
Under Assumption [ we demonstrate that the learned mapping function swap-label equivalent to the ground truth mapping (Theorem []. This equivalence ensures that while the specific labels assigned to meta states may differ between learned and ground truth mappings, the underlying causal relationships remain preserved. The proof demonstrates that any departure from this equivalence leads to higher LI norm of causal skeleton matrices and thus contradicts the optimality principle underlying our framework.
Swap-L.beel Equivalence:
Observational Equivalence: For overparameterized scenarios where the learned model contains more meta states than the ground truth, we prove that observational equivalence is maintained (Theorem 2). This result demonstrates that introducing additional meta states does not compromise the recovery of true causal relationships, since the learned mapping consistently preserves the same causal skeleton matrices across all states.
Intervention Reachability: We formalize the mathematical framework for determining which in- terventions are feasible given environmental constraints. Through matrix representations of state transitions and intervention capabilities, we establish conditions under which specific causal relation- ships can be identified through sequential interventions (Theorem 3 guidance for experimental design in causal discovery.
This analysis provides practical
Structural Assumptions:
Our results rely Assumption I which posits that learning from mixed datasets yields causal graphs encompassing the union of all contributing causal relationships. This assumption is both theoretically justified and practically reasonable, it ensures comprehensive coverage of causal dependencies present in heterogeneous data sources.
The theoretical framework presented in this section provides rigorous guarantees for the identifiability of causal structures while accommodating practical constraints inherent in real-world applications, thereby establishing the soundness of our proposed methodology.
A.2 Swap-Label Equivalence
In the context of Meta-Causal Graphs, swap-label equivalence captures the invariance of causal structure under permutation of meta state labels. This phenomenon arises when different labelings of meta states yield identical causal relationships across the state space.
Consider a Meta-Causal Graph G with meta states {ui, [2) and corresponding causal subgraphs (Mu, Mw as illustrated Figure 6a An alternative Meta-Causal Graph & with meta states {ū1, (2} and subgraphs {Ma, Ma, } is shown in Figure 6b
Two mappings C: : X + U and Ò X + i are swap-label equivalent if there exists a bijection g g: U + such that for any state a E X. the causal skeleton matrices satisfy:
Mo(w) My((())) Mo(a)
This equivalence ensures that the learned mapping Ć' preserves the same causal structure as the ground truth mapping C, despite potentially different meta state assignments. As illustrated Figure 6 such permutations preserve the fundamental causal relationships within the Meta-Causal Graph.
23

---

1 0 1
Mu,
1
x
Ru.
1 1
:)
1/
1
10
((x)= th
((0)
= =0
((x)= =
Mus
11 0
Ma.
(1 0
(1
0 ))
1
(b) Label-swapped meta states (@1, (u)) with corresponding subgraphs Ma, and Mag- Figure 6: Swap-label equivalence Meta-Causal Graphs. Left: Ground truth Meta-Causal Graph with meta states {#1, "2) and corresponding causal subgraphs (M Muz }. Right: Learned Meta- Causal Graph with meta states {ū) úg}} and subgraphs {Ma, Ma }. Despite different meta state labels, both mappings produce identical causal skeleton matrices for any state x E X, demonstrating structural equivalence under label permutation.
(a) Two meta states (u1, uz) with conresponding subgraphs Ma; and Mag.
A.3 Proof of Theorem
Proof. Consider the reconstruction objective for learning causal skeleton matrices from mixed data. For any clustering É: X' Ù, let Mr denote the learned skeleton matrix for meta state ú. Under Assumption I this matrix satisfies:
Me =1 L==((===
U
Now consider two cases:
Case 1: Ô' perfectly aligns with C' up to label permutation. Then there exists bijection U + U such that ((z) (((()) for all æ E X. In this case, Mg(x)) = M,, for each u E U, achieving perfect reconstruction of individual causal structures.
Case 2: Č merges states from different true states. Then there exist 21, 22 X with C(zi)
W1. # u2 C(x22) but ((11) = (x2)) û. By Assumption
Mr = Ma, UM
where U denotes element- wise logical OR. Since M., * Muz, we have Ma * Ma, and Ma * Maz-
Let || l| denote the L1 norm. For binary matrices, || equals the number of edges in the corresponding causal graph. Since the union operation can only add edges (never remove them), we have:
Mall 1| MM Mayll= ma II,, [| 111)
Moreover, since Ma, f Mug. there exists least one position (i,j) where Mu, [i.11 + Mug [i,
Mall > ||M,, and || ||Muzll
Consider the expected L1 norm under each partition. For the true partition C:
Exsx [|lMe[])]]]]
Σ P(C(a) u) M|MMlla NEU
For the suboptimal partition that merges distinct meta states:
Ex~x III
- = û) -I
set
Since merging increases the Ll norm (as shown above), and each merged cluster has || || > |Mul] for the constituent true meta states u, have:
Ez~x ((e))) [la] > E (Mlxx()) Iha]
24

---

0 0 (1
May 1
Ro. (
R
(G)) u,
(00) 0
060 R
May
0G0 0
Rs.
(a) Two meta states: partition (left) and Meta-Causal Graph (right).
(b) Three meta states: partition (left) and Causal Graph (right).
Figure 7: Observational equivalence under overparameterizatton. Left: Ground truth Meta-Causal Graph with meta states {U1 U2) and causal subgraphs (Muy Muz] Right: Overparameterized Meta- Causal Graph with three meta states (ū1, ú2 úz) and subgraphs (Ma, Mag, Mag.- Despite the additional meta state, both mappings produce identical causal skeleton matrices for any state E X, demonstrating that overparamcterization preserves causal structure.
A.4 Observational Equivalence
Observational equivalence captures scenarios where learned meta state mappings preserve causal structure despite overparameterization. Formally, two mappings C X U and C X Ü are observationally equivalent if for every state x E X: Mc(x) Mo(ay) where Mc(w) and M Mo(s) are the causal skeleton matrices corresponding to the assigned meta states. This equivalence ensures that the learned mapping Ò' captures identical causal relationships as the ground truth mapping C, even when the learned meta state space Ü is larger than the true space U. Overparameterization may result in redundant meta states, but the essential causal structure remains preserved.
Figure iillustrates this concept. Left: Ground truth Meta-Causal Graph with meta states (u,, u2 and subgraphs {Mar M. Right: Learned overparameterized Meta-Causal Graph with three meta states {ūn, úz ûz} and corresponding subgraphs (Ma, Ma, Mag Despite having additional meta state, both mappings yield identical causal skeleton matrices for any given state x E X, demonstrating observational equivalence under overparameterizatton..
A.5 Proof of Theorem 2
Proof. We prove that despite overparameterrzzaiin,, the learned mapping preserves the causal structure of the ground truth mapping.
Let Unsive {a E Ú Ex E X' s.t. C(x) = } denote the set of meta states that are actually assigned some state X.
Step 1: We first establish that the learned mapping cannot merge distinct true meta states. Suppose for contradiction that there exist 1 22 X' such that:
C(a1) = + (222) but ((z))=((az = à
By Assumption [ the learned causal skeleton matrix would be:
Ma = Mu, U Muz
Since Ma, A Mug. we have ||aal] 1/ > max M lh, [Maglli). increasing the expected structural
complexity as shown in Theorem T
This contradicts optimal clustering.
Step 2: Since distinct true meta states cannot be merged, each active learned meta state á E Ú Úsctive corresponds to exactly true meta state. That is, for each à Üseive there exists a unique u U such that: (a: ((x) û) C (x: C(x) = w.}
Step 3: For any X, (xx) and the corresponding true meta state is then by Assumption
Mr MM-
Therefore, Mc(e) MMM for all X, establishing observational equivalence.
25

---

Co0
0
Aco al con-a)
Ma.
Bi,
1
:
by
Ma,
1
Ma
1
1
800
000 -0.
11
0
o
Map
0
Ra.
9)
1
Ground truth: graph (left) and skeleton matrix (right). Figure 8: Illustration of Theorem 2 Each pair places the graph (left) alongside its causal skeleton
(b) Over-parameterrzed: learned graph (left) and skeleton matrix (right).
matrix (right): (a) ground truth; (b) overparameterized solution.
A.6 Disscussion Assumption
Let MG = [ x be a Meta-Causer Graph with K causal subgraphs corresponding distinct meta states u E: U. Consider dataset D where each sample is generated from of these causal subgraphs. Let SD S (1,2, K} denote the indices of causal subgraphs that actually contributed samples to D. When learning single causal graph & from the mixed dataset 'D (treating all samples as if they were generated from single causal subgraph), the estimated parent set for each variable Xj in & satisfies:
Assumption 2.
(X)(
U Pag, (X₃) SESp
Vj E [p).
Assumption D states that learning from mixed data generated by multiple causal subgraphs yields a unified graph capturing the union of all parent-child relationships across contributing subgraphs.
Figure 10 illustrates this assumption. The left panel shows the causal skeleton matrices M,, of individual subgraphs that contribute data to D. The right panel shows the learned causal skeleton matrix M from the pooled dataset. Red elements highlight the union property where edges present in any contributing subgraph appear the learned graph.
Mathematically, this relationship can be expressed as:
[i, =I
U Ma =
vi.j [p].
LuESS
where 1[-] is the indicator function, ensuring that an edge (i.j) exists in the learned graph if and only it exists in least one contributing subgraph.
This assumption is well-motivated: when data from multiple causal mechanisms are pooled without knowledge of their source, standard causal discovery algorithms tend to include all statistically supported edges to avoid missing true causal relationships. Violating this assumption would imply that the learning algorithm systematically ignores genuine causal relationships present in the data, leading to incomplete and potentially misleading causal models.
Assumption I shows that the LI norm of the learned causal skeleton matrix is larger than the LI norm the causal skeleton matrix of the contributing graphs. This also provides a way to show that the underparamettrrizaton of meta states will lead to the failure of learning the causal graph.
Figure shows an example of underparameteriiztton.. The upper figure shows the ground truth causal graph with three meta states and the corresponding causal skeleton matrix. The lower figure shows the learned graph with two meta states and the corresponding causal skeleton matrix. The underparameterizztion of the learned graph leads to the learned causal skeleton matrix different from the ground truth causal skeleton matrix (Equation 3)
However, overparameterization does not lead to the failure of learning the causal graph. Figure g shows example of overparamettrizatiin. ferent meta states to states which are generated from the same causal graph, the data learn each causal graph is still generated from the causal graph.
Although we may assign dif-
26

---

01
Mu, 1
M
(00) un
0
8oo
(00 mm m
Mag
0
((k)
Ma.
Ma.
1
0
causal graph (Ieft) and learned (a) Three graph (right).
(b) Two meta states: causal graph (left) and learned graph (right).
Figure 9: Illustration of the Assumption 1 The left figure shows the ground truth causal graph with three meta states and the corresponding causal skeleton matrix. The right figure shows the learned graph with two meta states and the corresponding causal skeleton matrix.
A.7 Proof of Proposition
I
1 0 1
0
Proof. We prove that the given intervention con- dition ensures unique identifiability by showing that no two distinct DAGs can be intervention- ally Markov equivalent under I.
Mu,
0
0
1
0 0)
Ma
1 0
1
1
11 0
0 1
0
Let Dr denote the interventional Markov equiv- alence class (I-MEC) of DAGs that are interven- tionally Markov equivalent D given I.
Muz
0
0
1 1
Step 1: Consider any edge - j the true causal graph D. By assumption, there exists I such that exactly one of [i.j} is !.
Figure 10: Illustration of the Assumption The left side shows the causal skeleton matrices of causal graphs which contribute the dataset for generating the data. The right side shows the learned causal skeleton matrices of the causal graph from the dataset. The elements in red high- light the differences between them.
Case 1: If i E I and & I, then in the inter- vention graph D(O, all edges into node are removed, but the edge i + (if it exists) re- mains. Consider any DAG D' with the reversed edge + i. In D'(",, this edge would be re- moved since E I, creating different interven- tion graphs: D() * (((1)
Case 2: j E and & I, then D all edges into j (including i + j) are removed. DAG D' with edge
would retain this edge in (() since é 7, again yielding D(I) + D)())
Step 2: Since every edge +
j E D satisfies the intervention condition, every edge direction uniquely determined by the interventional data, ensuring that for any DAG D' DI, we have D = ((( for all I I.
The intervention condition ensures that any edge orientation different from that in D would create distinguishable intervention graphs, contradicting interventional Markov equivalence. Therefore, [Dz] 1, meaning is uniquely identifiable.
Step 3:
Remark 1. Score-based causal discovery methods identify structures only up to Markov equiva- lence class, since different directed graphs can encode the same set of conditional independencies. Interventions, however, can break this equivalence: by actively perturbing variables, can dis- tinguish among graphs within the same class. For instance, in a causal chain A + - C, an intervention A affects C only through B. If B is not held fixed, a purely observational method might incorrectly infer a direct edge A - C. To fully recover the underfying structure, multiple interventions, potentially on more than variable, are required.
our frumework, this issue addressed two complementary steps. First, using interventional comparison (Proposition 7. we detect edge existence by contrasting causal subgraphs obtained from distinct interventions. Second, within each intervention subgruph, we apply a score-based refinement guided by structural minimality (Theorem Ty. This ensures that adding spurious edges
27

---

such as A + C would increase model complexity without improving fit, allowing the true causal structure to be uniquely identified.
A.8
Reachability of Interventions
Environmental constraints often prevent direct manipulation of every state variable in practice. We consider subset S. S [p] of variables permitting direct intervention, while others require multi- step sequences or remain unreachable. We define intervention reuchubility: state is reachable if attainable from the current state via finite allowable intervention sequences.
Each state variable X; takes discrete values, represented using one-hot encoding as z E {0, 10,1)% where exactly one element equals 1. The complete system state is z zi Bg where & denotes the Kronecker product and N = II Ti.
Zp {0. 1)
Remark 2.
The Kronecker product of one-hot vectors remains one-hot. 1f E {0, 1 and 22 E
(0, 1) are one-hot vectors, then 21 Z2 {0, 1 mm is also one-hot, with z a/nr+tu zı [r] 22 [v]
for r E [m], v E [n]. Since 21 and 22 are one-hot encoded vectors, we have zi(r] [r] E {0, 1 and 22[0] E {0. Therefore, uJ {0,1) and there exists only one ner + v such that z(nr + v] =
Thus, 2 is a one-hot encoded vector. Therefore, the Kronecker product of one-hot encoded vectors is still a one-hot encoded vector.
The Kronecker product is a one-hot encoded vector, which can represent the state of the system.
We define two key matrices:
Intervention Matrix F E (0, 1)NN; F[i,j] 1 if state j can be directly reached from state through intervention on variables Sc.
Transition Matrix T {0, 1xNN, T[i if the system naturally transitions from intervened state to state j.
Given a state vector {0, 13,, the intervention operation z' Fz produces:
N Flk, [zz[1]
z' z'[k
Since z is one-hot, z'[A] 0 if and only if there exists an index such that z and F[k, i] 1. This means state is directly reachable from the current state through intervention. The non-zero elements of Fz thus indicate all states achievable by single intervention step.
Similarly, the transition operation z" = Tz' yields:
2"kk] Σ
Thk [1
Here, z"kk] indicates that state i can be reached from some intervened state i where z z'[i] and lik i| 1. This captures the natural system dynamics following intervention.
Combining these operations, the complete system evolution under intervention is:
Z++1
TFz.
This composition first applies interventions (via F) to determine immediately accessible states, then applies system dynamics (via 7') find the resulting states after natural transitions. The non-zero elements in (TF))> zo represent all states reachable within intervention-transiii cycles from Zo- Theorem (Reachability Analysis). A state corresponding to index is reachable from initial state Zo and only if there exists k such that [(F))* zo 0.
Proof. We prove both directions of the equivalence.
28

---

(+); If state i is reachable from then by definition there exists a finite sequence of intervention- transition steps that leads from Zg state where the i-th component is active. This sequence corresponds to some power & of the operator TF, hence [(FF)*Z)] [ 0.
(+): ((TF)) > for some 2 O, then by the definition of matrix-vector multiplication, there exists computational path through k iterations of intervention-transiiie operations that activates the i-th state component, demonstrating reachability.
The reachability analysis has direct consequences for causal structure learning:
Corollary 1 (Intervention Feasibility). The state corresponding index i is feasible from initial state 2o by intervention if and only if there exists 2 such that [F (F(FF)zz zold]] 0.
Proof. This follows directly from Theorem 3 Here, F(TF)*z, represents states accessible for intervention after & cycles, where the leading F captures the final intervention step.
This constraint fundamentally limits causal discovery scope in practical settings. The Curious Muusality-Seeking Agent must therefore operate within feasible intervention constraints while maxi- mizing causal structure identification. We provide an example to illustrate these results.
Case 1. Consider a system with two binary variables D2, yielding four possible states as shown in Table 3 As 21 and 22 one-hot encoded vectors, the Kronecker product of 21 and 22 is
Table 3: One-hot encoding of æ and and their Kronecker product.
31
F2
21
2/
2
21 X 22
0
1,0
1,0
0,
0
0,
1]
1,0 T
0, 0,,,0,0
0
1,0 T
1
1,00,1,0
1
00,
1]
1]
10,
4-dimensional one-hot encoded vector, which can represent the state of the system. The space of the system be represented 2 {[1,0, 0, 0]", ,0, 1, 0, O, [0, 0, [0, D. O. T. For 24, E 2, if 2j can be reached 2; by intervening on Ie we can denote this as Fii,j= =1
Intervention Constraints: Suppose only can be directly intervened. The intervention matrix F allows transitions between states that differ only in ø1;
0 0
F =
1
System Dynamics: After intervention, assume the system exchanges the values of xı and x2. The transition matrix becomes:
1
0
0
07
0
T
0
1
0
0
1
Reachability Analysis:
Starting from 0]
(state (0, 0))
Step 1: Direct intervention possibilities:
Fzo
States I and 2 are accessible for intervention.
29

---

Step 2: After one intervention-transittoo cycle:
T
The system can reach states I and 3.
Step 3: Interventions possible from these new states:
F(TFZo) = F
All states become reachable for intervention within two cycles, demonstrating that despite initial constraints, the system's 's dynamics enable full state space exploration.
A.9 Action Intervention
Our method identifies context-dependent causal graphs (meta-graphs) via interventions. Actions are how the curiosity-driven agent realizes those interventions. When actions directly intervene each state (e.g., chemical environment), we deploy targeted interventions to isolate edges and quickly distinguish meta states and their subgraphs. When not all interventions are feasible, we rely on the reachability assumption.
Verification Results. We add experiment based manipulation task to demonstrate that our method can handle cases where actions cannot intervene on all variables. We also modify this environment test generalization. The results show that method maintains strong performance and generalizes well, even when only a subset of variables is directly intervenable.
Table 4: Performance on partial-intervention manipulation tasks.
Reward
Small Magnetic Force (x0.02)
High Ball Density (x10)
Extra Table Friction
MLP
4.011 = 0.030
3.999 = 0.038
3.886 + 0.241
FCDL
4.482 = 0.622
4.451 0.657
4.193 0.655
MCG
6.173 = 0.255
5.083 = 0.245
5.517 + 0.134
A.10
Lossy Representation and Misclassification Probability
We analyze how lossy representation affects the probability of mapping samples to prototype embed- dings. Table 5 summarizes the key notations used in this derivation.
We assume that for each state r with true meta state C(x) = u, the probability that is mapped by the encoder to the codebook entry index ûg is P = P(Ć(z) = ås C(x) = u). This probability reflects the representation power and discriminability of the encoder, indicating how likely samples of the same true meta state are mapped to different codebook entry.
From Definition (xx) and C(x) are observationally equivalent if, for all un E Ô and all 01,02 E {o Ō(x) 0 J. it holds that C(æz) = C(zz).
The probability that C(x) and C(z) are observationally equivalent is given by P(C(x) u, Vx' x ((x) = úμ C(z') = u) = (1 Pk +
where n is the number of samples.
Hence, the overall probability of observational equivalence is
P Pequiv
Σ Pk
Σ
pal Pk + pM
uEU
30

---

Table Notation summary.
Symbol
Meaning
U u
Ground-truth meta states
Ü us
Prototype (codebook) embeddings
C(x)
True meta state of sample xr
C(x)
Codebook entry selected by encoder
Hu Pk Dueu MuPk
Prior probability of state z with true meta state a Probability of mapping sample to codebook entry cū2
The probability that all samples assigned to same codebook entry index ((x) belong to different true meta states is
Paint ΣΣ (ps - (1PP)) - Pk: + Pk
NEU
The probability of misclassification, i.e., that codebook entry corresponds mixed true meta states, is
P misclass 1
Σ Pk E1- (1-
Considering lossy representation, the misclassification probability simplifies to
P(misclassification) =
-E-E
(1pP(( - Pk + pkpu))
WEU
If nppp + ph L, we approximate
P(misslassification)
phpm)) Σ Pk E(uPpE)
In a more accurate representation, for each codebook entry index kt the distribution of PaP concen- trates on a single true state one ph increases while others decrease, thereby increasing Σ,, (MuP)) and reducing each term j >,, (MmPi) Consequently, higher representation accuracy mono- tonically decreases the overall misclassification probability, reaching zero when each prototype corresponds purely to one true meta state.
Experimental Details
B.1 Environment
Chemical. We evaluate our method's performance on learning context-dependent causal structures using the Chemical environment [39]. This environment consists of 10 objects, each capable of taking one color states. An action selects a target object and changes its state, triggering cascading changes to all dependent objects according to the underlying causal graph. The objective is match each object's color to specified target configuration.
The Chemical environment features multiple causal structures that switch dynamically based on the system state, making it ideal for evaluating Meta-Causal Graph learning. We consider two experimental settings in this environment:
Full-Fork (Fork): The causal structure alternates between two graphs depending on the root node's color. When the root node is red, fork structure is active; otherwise, cully-connected structure governs the system dynamics.
Full-Chain (Chain): The causal structure alternates between full and chain configurations. A red root node activates the chain structure, while other colors trigger the fully-connected structure.
These settings test our agent's ability to: (1) discover the latent meta states (root node colors)
that determine causal structure transitions, (2) learn the corresponding causal subgraphs through
31

---

interventional exploration, and (3) generalize to unseen state configurations during evaluation. During the test, some nodes are corrupted with noise. The agent needs to learn the causal graph to match the colors of nodes to the target. The agent starts from random color configuration and must transform a 10-node graph to match target color pattern, Actions intervene on nodes, changing their color and that of all their descendants according to the hidden causal graph. The episode reward is the negative Hamming distance to the target (0 for perfect match).
Magnetic. The Magnetic environment is robot arm manipulation task built on the Robosuite framework (Figure 4) The environment contains two objects: fixed box and a movable ball, whose colors indicate their magnetic properties. When both objects are red, they exhibit magnetic attraction, causing the ball's trajectory to be influenced by the box's position. The magnetic properties of each object are randomly assigned at the beginning of each episode.
B.2 Reinforcement Learning Algorithm
For fair comparison, we use the same model based reinforcement learning algorithm for all the baselines and our method. We use the cross-entropy method (CEM) [1] to sample the action based on the predicted transition dynamics. The detailed hypper-parameten are shown in the Table 6
Table 6: CEM parameter.
Chemical
CEM parameters
Magnetic
full-fork
full-chain
Planning length
3
3
Number of candidates
64
64
64
Number of top candidates
32
32
32
Number of iterations
5
5
Exploration noise
N/A
N/A
le-4
Exploration probability
0.05
0.05
N/A
Action type
Discrete
Discrete
Continous
B.3 Environment Configurations
The detailed environment configurations are shown in Table 7
Table 7: Environment configurations
Chemical
Paramters
Magnetic
Fork
Chain
Training step
1.5 x 10
1.5 10a
x 10
Optimizer
Adam
Adam
Adam
Learning rate
le-4
le-4
le-4
Batch size
256
256
256
Initial step
1000
1000
1500
Max episode length
25
25
25
Action type
Discrete
Discrete
Continous
B.4 Partial Observability
Table 8: Prediction accuracy under partial observability the Chemical environment.
Prediction Accuracy
n
n 6
MLP
31.93 + 0.16
32.47
1.61
33.72 + 2.11
FCDL
63.75 + 13.75
53.89
11.95
52.11 + 16.43
MCG (ours)
73.52 + 9.56
63.71
5.79
57.05 + 14.63
32

---

0.00
5:00
1:59
9:00
8.00
0.00
5.00
dg der my d5 oue ogy Ck
0.00
8:00
9:00
8.00
0/9
6.00
120
0.08
0.03
8.00
994 992
2.00
8:0
9:00
5.00
0.00
1.20
A00
0.00
5.00
oet o4t 04 04 044
0:00
8:00
0N 0g3
1.
0:00
1.00
L:0
2:20
1gg OM
LG0
621
9.04
0.00
0.00
LAX
100
120
L00
-02
0:
1:c0
3:00
Can
0:00
0:00
1:2
100
1.00
1:00
1:00
(a) Causal subgraph of Full.
0.05
5.50
LAD
-04
AUD
0:08
0.50
5.00
0.00
0.00
-04
1:30
0:00
8:00
0.00
5-10
04S
X1X
PM
6:00
0.80
1:00
PXT
8:00
noe 0.5 135
5.00
-02
-.2
0.0
6:
(b) Causal subgraph of Fork.
3,00
0.00
8:00
0.00
9.00
JU00
GUR
0.00
8.00
0.0
LOP
LGO
3:00
0.00 8.50
DAN
220
300
0.00
8:00
80
5:00
0:00
5:00
1.59
0.00
1,09
dee
-04
ogt
oxk
ADT
6.00
1:00
L.0
6:0 aot
120
6X A09
(c)
Causal subgraph of Chain.
Figure 11: Comparison of learned causal graphs in task Chemical. MCG (ours) is the proposed method and the reference is the causal graph learned by FCDL.
33

---

B.5
Extended Results
visualize the learned causal subgraphs in Figure [T]
The results demonstrate that our method
effectively captures the underlying causal structures of the environment.
B.6 Compute Resources and Environment Details
Most experiments conducted server equipped with AMD EPYC 64-Core Processor
(24 physical cores), supporting 32-bit and 64-bit modes, with 96 MiB L3 cache and 12 MiB L2 cache. The machine equipped with an NVIDIA A100 PCIe GPU with 80GB memory (driver version 575.51.00 .03, CUDA version 12.9).
B.7 Code and Demo
The code available at and demonstrations be found at http::///sssseegeeeee//////////a
B.8
Prediction accuracy for three intrinsic terms
Table Prediction accuracy for three intrinsic terms. All values are reported in percentage %.
Intrinsic Term
1=2
n=6
n=2
1=1
n=6
Prodiction Uncertainty
62. 02.25-2 + 16.73
59.46 22.38
52.25 + 28.3223888 16
59.52
8.82
51.81 = 5:47
59.28
5.03
Feature Discrepancy
46.11
41.17
43.67
2.47
56.01
13.13
51.88
7.22
58.27
8.53
Predictive Distribation Discrepancy
54.23 6.72
48.27 9.74
43.92
1.85
3.11 36.83
40.33 4.89
45,40
2.83
Edge Entropy
63.18:13.94
50.4719.87
50.0418.56
51.9911.588
49.7824.11
49.69:15114
B.9 Efficiency
We compare other methods (GNN and MI.P) with longer training rounds and more parameters to show efficiency of our method. Results are reported in Table T0 and Table
Table 10: Comparison under the same model size. We evaluate the efficiency of our method by
comparing MCG with GNN and MLP trained for longer rounds and with larger parameter counts.
Acc n=2
Acc n=4
Acc n=6
Reward n=2
Reward n=4
Reward n=
GNN (30k round)
38.13
25.81
20.75
6.83
6.06
6.83
GNN (25k round)
36.37
26.55
20.73
6.84
5.98
6.19
MLP (30k round)
30.47
28.94
27.84
6,20
6.13
6.71
MLP (25k round)
31.47
30.32
30.42
6.18
6.03
6.23
MCG (10k round)
47.01
48.37
46.35
11.96
10.83
10.70
MCG (12k round)
48.30
48.40
47.28
12.45
10.77
10.84
MCG (15k round)
71.56
59.17
49.58
14.65
14.06
13.28
Table 11: Comparison under the same number of training rounds. MCG is compared with GNN and
MLP models trained for the same number of rounds but with different model sizes (parameter counts shown in parentheses).
n=2
Acc n=4
Ace n=6
Reward n=2
Reward n=4
Reward
MI.P (parameter x3.71)
24.65
25.05
23.71
6.29
6.50
7.93
GNN (parameter x1.66)
37.58
30.61
30.05
6.34
6.49
7.90
MCG (parameter x1)
71.56
59.17
49.58
14.65
14.06
13.28
34

---

C Extended Related Work
C.1 World Models
Contemporary research on world models can be delineated along two principal trajectories, each characterized by fundamentally divergent objectives. The first research trajectory conceptualizes world models as instrumental components, primarily serving either as predictive mechanisms to facilitate planning processes
training apparatuses for policy optimization. Conversely, the second research trajectory approaches world models as generative frameworks, with the explicit objective of predicting future environmental states with high fidelity.
Within the first trajectory, numerous approaches leverage world models as predictive mechanisms to facilitate planning processes. Notable exemplars include methodologies employing Monte Carlo Tree Search (MCTS) to identify optimal action sequences [65] 12 and techniques utilizing cross-
entropy methods efficiently sample continuous actions 225 Complementary to these, substantial corpus of research utilizes world models to augment policy learning, wherein the learned dynamics models either provide supplementary supervision signals function synthetic environments for policy optimization. This subcategory encompasses seminal works such as Dreamer [2 23
241. SimPLe (38], IRIS (53), AIIRIS [54], and DART (]]. These methodologies predominantly employ model-based reinforcement learning frameworks, wherein learned dynamics models generate synthetic data facilitate policy model training.
The second trajectory is predominantly focused on the development of sophisticated generative models designed to predict future environmental states with high verisimilitude. Representative examples include DIAMOND [3, Navigation world models [.. Oasis [[]], and The Matrix [[6]. These models engineered to capture the underlying stochastic dynamics of complex environments and generate realistic future states conditioned on current states and selected actions.
Notwithstanding these advancements, it is imperative to acknowledge that methodologies relying predominantly on statistical correlations frequently exhibit performance degradation when confronted with distributional shifts in environmental conditions, thereby compromising their capacity for robust generalization across diverse scenarios.
C.2
Causal Discovery for World Models
Causal discovery offers a rigorous analytical framework for addressing the inherent challenges associated with distributional shifts in world models. By elucidating and exploiting causal relation- ships rather than mere statistical correlations, these methodologies significantly enhance both the interpretability and generalization capabilities of learned models, thereby facilitating more robust Eecision-mmaing processes in complex, stationttooy environments.
The methodological landscape of causal discovery can be taxonomized into two principal categories: constraint- approaches and score-based approaches. Constraint-basee methodologies, exempli- fied by the PC algorithm [68], employ conditional independence tests to systematically infer causal relationships among variables. Conversely, score- approaches utilize statistical evaluation metrics to assess the plausibility of alternative causal structures. Prominent instantiations include the Greedy Equivalence Search (GES) [51]. which implements a greedy search algorithm to optimize predefined scoring function, and methods leveraging the Bayesian Information Criterion (BIC). Both methodological paradigms endeavor to recover causal graphs by identifying the Markov equivalence class of the underlying causal structure. Nevertheless, it is imperative to acknowledge that the causal graph learned through these approaches is not uniquely identifiable. as multiple distinct causal architectures can manifest identical conditional independence relationships.
To mitigate this fundamental identifiability challenge, contemporary advancements in causal discov- ery have increasingly focused on incorporating interventional data to enhance the discriminability of causal structures (69 52 33
66 33). Despite these methodological innovations, significant limitation persists: these approaches frequently operate under the restrictive assumption that the underlying causal graph remains temporally invariant throughout the learning process. This station- arity assumption may be violated in dynamic environments, wherein causal structures can evolve temporally due to myriad factors.
72
6
35

---

C.3 Comparison with FCDL
Motivation and Comparison to Context-Dependert Causality. The initial motivation for our work came from a fundamental question: Why do causal discovery methods often fail in open-ended environments?
We observed that causal models aim to describe the rules of the world. However, in open-ended settings (e.g., open-ended games, multi-agent systems, or LLM-based dialogue), no single observation window can capture all possible events. This leads to instability: a causal relation valid in one context may fail in another. For example, Newtonian physics holds at the macroscopic level but breaks down the quantum scale. FCDL also emphasizes this issue (but it's not the first one), arguing that complete causal law must account for variation across different contexts. Therefore, our high-level motivation aligns with them. many previous works have explored this motivation context-dependent causal structures.
However, FCDL is not the first one discusses about context-dependent causality. The idea of context- dependent causality was formally introduced by earlier Huang et al. [29]. and also discussed earliest in Chapter 10 of Pearl's Causality [56].
Exploration and the Limitation of Passive Learning. A key challenge in scaling up to open-ended worlds is how explore effectively. Relying solely on random exploration limited observed data insufficient to discover hidden causal mechanisms. For instance, in scientific discovery, new causal knowledge often emerges only after targeted investigation of paradoxes. Similarly, in dialogue agents, causal relationships may even reverse depending on the context. Unlike prior works including FCDL and Huang al. [29], which rely solely on observation passive data, our method explicitly considers the unobserved causal space, and how newly explored data may lead to shifts in both the causal graph and the context boundary.
Causality-Seekingg Agent. We aims to go beyond passive modeling to active exploration for causal discovery. Our curiosity and intervention-based strategy address core limitation of data-driven methods: the inability to discover unknown causal structures. Our primary theoretical contribution the construction of formal intervention-batt meta-causal graph, with associated theorems and learning framework.
C.4
Connections to Model-Based RL and Probabilistic Graphical Models
Our framework can be viewed as a special case of model-based RL (MBRL): we learn a world model and use it to plan. The difference lies in the causal structure learn and use. Classical factored or relational MBRL methods learn structured predictive models 20 42 600, which improve sample-effiiinncy but typically remain correlational. By contrast, explicitly model Meta-causal graph whose edges carry interventional semantics (actions are treated as interventions). and we pair structure learning with intervention design and identifiability analysis for context-dependent mechanisms (meta states).
From the viewpoint of probabilistic graphical models (PGMs), our method instantiates structural causal model (SCM) with a discrete context variable u (the meta state): each context selects a graph G. and its causal skeleton My. while the decoder D outputs a probabilistic mask M, that we discretize via Gumbel reparameterization. This preserves the graphical factorization of PGM but endows with interventional semantics in the sense of Pearl's do-calculus [56]. Practically, this differs from standard PGM-based world models by (i) targeting causal edges rather than solely predictive factors, (ii) using auriosity-driven interventions to break equivalence classes, and (iii) handling changing mechanisms via meta states with adaptive, minimal codebook.
Finally, our curiosity module is compatible with latent-variable exploration used in POMDP-style MBRL: intrinsic objectives (entropy, feature discrepancy, predictive-NI.L. accuracy) can be plugged in as information-seeking criteria, but here they are aimed al reducing uncertainty over causal structure (meta states and edges), not only over hidden state trajectories.
36

---

D Theoretical Connections between Meta-Causal Graphs, Open-Endness, and Gödel Machines
Open-endedness is defined in terms of continuous novelty and learnability, highlighting its significance for creating artificial superhuman intelligence (ASI) 30). Additionally, the Gödel Machine concept proposed encapsulates self-referential improvement mechanisms through formal proof-driven code
rewriting. Here, we analyze the theoretical relationships between the MCG framework, open- endedness, and the Gödel Machine [63].
D.1 Open-Endness and Meta-Causal Graphs
The Meta- MettaCausaal Graph framework establishes connections with open-endedness and self-improving systems. Open-endedness is characterized by continuous generation of novelty within comprehensible boundaries, foundational requirement for advanced artificial intelligence systems. Our MCG framework demonstrates key open-ended properties in several fundamental ways:
First, the curious causality-seeking agent operationalizes open-ended exploration through curiosity- driven intervention strategy. By maximizing entropy-based reward signals focused on regions of causal uncertainty, the agent persistently generates novel interventions and discoveries. This mechanism directly addresses the novelty requirement of open- neen-endedness, as the agent autonomously uncovers and probes previously uncharted causal relationships and latent meta states.
The MCG framework maintains interpretability by organizing new causal knowledge within an explicit meta-causal graph. As mechanisms are discovered, they are systematically integrated, ensuring ongoing comprehensibility despite growing complexity. Vector quantization further enables efficient and semantically clear assignment of novel observations to distinct meta states.
The MCG framework thus marks step toward genuinely open-ended learning systems capable of autonomously exploring, understanding, and adapting to complex, dynamic environments with evolving causal structures.
D.2 Gödel Machines and Meta-Causal Graphs
It is important to note that while the MCG framework does not implement Gödel Machine-style code-level self-rewriting, its core mechanism nevertheless embodies form of self-improvement focused on its own knowledge structure (the Meta-Causal Graph). The curious causality eausalit--seeking agent, driven by curiosity-ban -base rewards, actively collects new data and, upon encountering prediction failures high uncertainty, dynamically refines and expands its set of causal subgraphs and meta state mappings. This process ensures the agent can continually update and improve its world model in open environments, achieving self-correction and knowledge-lewe self-evolution.
E Implementation Details
E.1
Meta-Causal Graph Learning from Experience
We learn causal subgraphs under latent meta states from agent experience. Following FCDL's pioneering use of the VQ-VAE architecture for causal-graph discovery [3I, we adopt the original VQ-VAE framework of den Oord al. (74] learn the causal subgraph for each meta state. Specifically, the process contains two steps: (1) identifying the latent meta state, which activates specific causal subgraph, and (2) learning the subgraph from the agent's explored experience data.
Identifying Latent Meta State: We embed the observed state and assign to a corresponding meta state E U by vector quantization as follows. We define the meta state assignment as C(x) = arg min Exx) - zu 1|-- where E: X + Rdd is an encoder mapping the observed state d-dimensional embedding, and (zu)nev Rd is learnable codebook of prototype embeddings, each representing a distinct meta state.
Subgraph Learning: After obtaining the embedding 2n for a given state, we employ a decoder network D Rd + [0, 1|xx to predict probability matrix M., which estimates the underlying causal skeleton matrix M,, of G.. Each entry [. denotes the probability that the directed edge
37

---

i ++ j is present in causal subgraph associated with meta state u. We then sample each Mu from Bernoulli( (M. ji) using the Gumbel-Softmax reparameterization [35 50), enabling end-to-end
gradient- learning of the discrete structure. To update the codebook, we use the following update rule:
Lquantization (x)) - 2u11 + B([sg(z,)) (x))
where sg(:) is the nop-gradient operator. The first term is the reconstruction loss and the second term is the commitment loss to avoid the output of the encoder growing arbitrarily (74].
38


Reviews:
[{"review_id": "ILkzUY8U2C", "reviewer": "Reviewer_KoPL", "summary": "This paper proposed a new definition of Meta-Causal Graph as a world model, and used curiosity-driven intervention policy for causal discovery. They conduct experiments in chemical and magnetic tasks to verify. \n\nFrom my perspective, the main contributions are two parts: 1. Use Meta-Causal Graph for causal discovery. 2. Design a curious reward fothe r exploration intervention. The proposed learning process is reasonable. But the experiments are not convincing.", "questions": "Please see the pros and cons.", "limitations": "Please see the pros and cons.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths: \n\n1. Theoretically, they derive sufficient conditions ensuring that strategically chosen families of interventions uniquely identify meta-causal graphs, overcoming observational limitations.\n\n2. They use meta meta-causal graph for causal discovery, and design a curious reward. The learning process is reasonable.\n\nWeaknesses:\n\n1. Need more experiments to verify the performance of the proposed method. Maybe other experimental environments can be conducted, not just chemical and manipulation. Moreover, you can revise the Magnetic tasks for generalization verification. The experiments in generalization are not enough.\n\n2. Clear the words for the method description. In lines 51-56 of the method description, the learning process is not described clearly. \n\n3. Explain the Meta-causal graph: from my understanding, the meta-causal graph is a combination of some graphs and their specific indicators. Right? For policy learning, learn those graphs and indicators, and then use them for downstream policy learning? \n\n4. Explain the world. The author mentions that the method is a meta-causal world. I think the model learned in this work is a dynamics model, like CDL, Grader. It is not a world model. \n\n5. More recent works should be discussed, for example:\n\n        [1] Wang Z, Wang C, Xiao X, et al. Building minimal and reusable causal state abstractions for reinforcement learning[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2024, 38(14): 15778-15786.\n\n        [2] Cao H, Feng F, Fang M, et al. Towards empowerment gain through causal structure learning in model-based reinforcement learning[C]//The Thirteenth International Conference on Learning Representations. 2025.\n\n        [3] Yu Z, Ruan J, Xing D. Learning causal dynamics models in object-oriented environments[J]. arXiv preprint arXiv:2405.12615, 2024.\n\nI hope that the author can explain the above questions.", "was_revised": true, "final_justification": "The authors' rebuttal addressed my concerns. I suggest conducting more experiments in more challenging environments (Meta-World/CausalWorld) to demonstrate that the proposed method is a genuine meta-cause-and-effect world in the future. I maintain my positive score.", "rebuttal": "We thank the reviewer for highlighting our innovative Meta-Causal Graph framework, the effective curiosity-driven exploration strategy, and the strong theoretical guarantees for identifiability.\n\n> **W1: More environments & experiments in generalization.**\n\nThank you for your suggestion. We **added new environment** and **adjusted magnetic force, density of ball and friction of the table** to test **the generalization ability** of our method. The results are shown below and the metric is epsiode reward on downstream tasks:\n\n|Reward↑|Small Magnetic Force ($\\times 0.02$)|High Ball Density ($\\times 10$)|Extra Table Friction|\n|-|-|-|-|\n|MLP|4.01|4.00|3.89|\n|FCDL|4.48|4.45|4.19|\n|**MCG**|**6.17**|**5.08**|**5.52**|\n\nWe also **added the new experiment** to test the generalization ability of our method under varying levels ($n$) of noise (line 265-268) **when some variables are not observed**. The results are shown below and the metric is prediction accuracy:\n\n|Acc↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|FCDL|60.10|52.97|45.80|46.37|**47.88**|46.62|\n|MLP|31.56|31.67|31.93|25.95|24.71|21.66|\n|**MCG**|**66.94**|**59.19**|**46.35**|**49.93**|45.41|**48.09**|\n\n> **W2: Improve clarity of the method part.**\n\nThank you for your suggestion. Here is a clearer summary of the learning process:\n- For each observed state, the agent predicts a **compact meta-causal graph (MCG)** representing the current context’s causal structure.\n- The agent selects actions based on a curiosity-driven reward, which encourages exploration of uncertain or unseen contexts.\n- The agent uses the learned MCG to model and **predict environment dynamics**.\n- As the agent collects new experience, both the meta-state assignment and the corresponding causal graphs are iteratively updated. \n\nWe will improve the clarity of the method description in the revised version.\n\n> **W3: Explanation of the Meta-causal graph.**\n\n- Yes, your understanding is correct. The meta-causal graph is a collection of causal subgraphs, each associated with a specific meta state indicator.\n- We use the meta-causal graph for downstream tasks, where the agent learns to **predict the transition dynamics based on the meta state and its corresponding causal subgraph**.\n- Importantly, these meta states are **unknown** to the agent and must be discovered through interaction, which makes the problem challenging. The joint identifiability of both the causal subgraphs and the underlying meta states requires rigorous theoretical justification, as it is a fundamentally difficult problem.\n\n> **W4: Why is the Meta-Causal Graph considered a world model?**\n\n- The definition of a world model remains a subject of ongoing debate [Ding et al., 2024]. Generally, **a world model is any system that models the environment’s dynamics**—i.e., it learns the function mapping (state, action) pairs to next states. In this broad sense, CDL and Grader are also world models. \n- Our approach aligns with the broader definition of a world model: we first infer a compact **meta-causal graph** from each input, capturing the context-specific causal structure, and then use this representation to **predict future dynamics**. \n- This enables our model not only to predict transitions, like a dynamics model, but also to maintain an **abstract and interpretable understanding of the environment’s underlying structure**.\n\n> **W5: More recent work should be discussed.**\n\n- We appreciate your suggestion and give a brief comparison here. We will include more recent works and discuss comparisons in the revised version.\n    - [Wang et al., 2024]: This observation-based, passive learning method **cannot actively explore or adapt in open-ended worlds**. In contrast, our intervention-based approach **actively discovers unseen meta-states to learn a Meta-Causal Graph**, making it well-suited for open-ended environments.\n    - [Cao et al., 2025]: Their causal structure is **global and fixed**, not adapting to state or meta-state changes. Our framework explicitly models **multiple context-dependent causal graphs, allowing discovery of diverse causal mechanisms** across meta-states.\n    - [Yu et al., 2024]: OOCDM learns a **shared causal structure** for each object class through passive observation, **without context-specific graphs or active exploration**. Our method discovers **multiple context-specific graphs** via **active, curiosity-driven interventions**, enhancing exploration and identifiability.\n\n- Additionally, We **added new baselines**: causal discovery using transformer attention weights (\"transformer\") and the Jacobian matrices of MLP layers (\"sandy-mixture\") from [Pitis et al., 2020] and [Urpí et al., 2024], for comparison. The results are shown below. The metrics used are next-state prediction accuracy and the average episode reward achieved by model-based planning on downstream tasks.\n\n|Acc↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|transformer|24.45|22.49|21.41|29.62|30.37|29.78|\n|sandy-mixure|31.93|32.47|33.72|30.08|29.31|27.43|\n|**MCG**|**63.18**|**50.47**|**50.04**|**51.99**|**49.78**|**49.69**|\n\n|Reward↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|transformer|6.20|6.13|6.71|6.67|6.73|7.40|\n|sandy-mixure|6.48|6.10|6.50|7.03|6.73|7.17|\n|**MCG**|**14.65**|**14.06**|**13.28**|**13.82**|**12.49**|**12.45**|\n\nPitis, Silviu, et al. Counterfactual data augmentation using locally factored dynamics. NeurIPS, 2020\n\nUrpí, Núria, et al. Causal action influence aware counterfactual data augmentation. arXiv, 2024.\n\nYang Liu, et al. Learning world models with identifiable factorization. NeurIPS, 2023.\n\nZizhao Wang, et al. Building minimal and reusable causal state abstractions for reinforcement learning. AAAI, 2024.\n\nHongye Cao, et al. Towards empowerment gain through causal structure learning in model-based RL. arXiv, 2025.\n\nYu Zhongwei, et al. Learning causal dynamics models in object-oriented environments. arXiv, 2024.\n\nJingtao Ding, et al. Understanding world or predicting future? a comprehensive survey of world models. ACM Computing Surveys, 2024."}, {"review_id": "OfEAvNqmsd", "reviewer": "Reviewer_hzVP", "summary": "This paper aims to discover context-specific causal structures by learning a meta-causal graph based on meta-state estimation and targeted interventions that trigger specific meta-states. The estimation process relies on learning a quantized latent meta-state space via VQ-VAE, while intervention targets are selected based on uncertainty measurements of the graph edges. Together with likelihood matching (i.e., dynamics prediction of future states) and sparsity constraints, the framework jointly learns these components in an iterative loop to discover context-specific causal graphs.\n\nIn the experiments, the authors demonstrate the effectiveness of the approach on simulated benchmarks with chain and fork causal structures, as well as on RobotSuite tasks involving different object interactions to validate its practical utility. Overall, the approach presents a clear and straightforward design. \n\nIn addition, the paper provides theoretical guarantees on the identifiability of meta-states and subgraphs under certain assumptions, along with an analysis of intervention reachability.\n\nIn general, I believe this paper makes meaningful contributions to the field of context-specific causal discovery. However, I have some concerns and questions, which I'll specify in the following sections", "questions": "As I have already listed most of my questions in the weaknesses section to keep them coherent and easier to address during the rebuttal, I will only raise two additional, more general questions here:\n\nQ1. How can the proposed context-specific causal discovery approach be effectively integrated into world models, as motivated in the introduction? More concretely, how can one efficiently learn and estimate changing causal structures in world model environments? Are there any additional structural assumptions needed? Furthermore, how should actions be incorporated into the framework? Is it possible to leverage this structure discovery as part of reward learning, for example, by treating it as an intrinsic reward for skill discovery?\n\nQ2. I have reviewed the theoretical analysis, and overall, the assumptions appear reasonable with no major issues. However, the theoretical guarantees seem to rely on the assumption of observational equivalence. That is, assuming access to a perfect generative model for learning dynamics, which ensures identifiability of meta-states and subgraphs. While I understand this is a common assumption in the causality literature, it offers limited direct guidance for practical methodological design. It would strengthen the work to provide more discussion on this connection. Additionally, in practice, the learned representation is often lossy to some extent. Have you considered further theoretical analysis on how such imperfect (lossy) representations may affect the identifiability of meta-states and the causal graph?\n\n- **final note**: as mentioned in the weaknesses and questions above, I feel that the current version of the paper does not yet meet the acceptance bar.  The main concern is the significant overlap in objectives and theoretical foundations with FCDL, also the experimental validation is lacking in terms of comparisons with other intrinsic motivation methodologies and more comprehensive benchmarks. \nThat said, I am optimistic about the direction and potential of this work and would like to see the authors' response. For points W1–W5, I believe further discussion and clear plans for revision are necessary. Additional experiments are welcome but not strictly required for all points, given the short rebuttal window. Based on the current version, I lean toward a reject, but I would be happy to raise my score if these concerns are adequately addressed and clarified.\n\n***Given the short rebuttal window, I would be happy to see: (1) a formal discussion clarifying the contributions beyond FCDL, particularly regarding the critical role of the curiosity term (see W1-W2) and the empirical attainability of the curiosity-driven objective; (2) additional evaluations, or at least a clear explanation of what is feasible to provide and what is not within the current scope (see W3); (3) anwers to other questions listed above.***", "limitations": "Yes, the authors provided a discussion on limitations.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 1, "strengths_and_weaknesses": "### **Strengths**\n\n(+) [about problem setting] The problem setting of context-specific causal discovery is an important and meaningful direction within the broader field of causal discovery. The setup in this paper, where the agent can actively intervene in the system, is interesting and potentially impactful though some prior works have explored similar ideas.\n\n(+) [methodology] The methodology is clear and straightforward, and the proposed objectives align well with the problem setting. These objectives could also potentially be useful for other tasks like MBRL or general world models. \n\n(+) [Presentation] The presentation is clear and easy to follow. \n\n### **Weaknesses**\n\n**W1**: [Overall technical contribution] The first weakness concerns the methodological novelty. To be clear, I generally do not favor raising this point (as I find excessive focus on \"lack of novelty\" unconstructive in many reviews), but in this case, the algorithmic framework appears to be heavily based on prior works: FCDL [Inwoo et al., 2024] and NCD [Inwoo et al., 2023]. I understand that this paper introduces a curiosity-based intervention module, but this primarily serves as a mechanism for selecting intervention targets. In my view, the core components of learning the meta-causal state representation and identifying subgraphs via VQ-VAE are identical to FCDL. That said, I do not consider this a critical issue that would warrant rejection. However, I believe the paper should be better structured to explicitly acknowledge this connection, for example, by clearly stating that the approach builds on FCDL and extends it with a curiosity-driven intervention selection strategy. Such clarification would likely require a revision of the paper. Also, maybe adding one table to compare them and expanding the related work part to discuss the details. \n\n**W2**: [Interventional design] As this could be the most critical module in this paper, I would suggest touching on a more detailed analysis on that part to make sure it is technically robust and correct. \n\n- The intervention selection is based on entropy. How do you ensure that this is tractable in practice? Specifically, what happens if the posterior distribution estimation is inaccurate during the early stages of learning?\n\n- Did you consider other commonly used curiosity-based measures for intervention selection, such as prediction error-based curiosity (Pathak et al., 2017), model disagreement (Kim et al., 2020),  or curiosity-driven replay (Kauvar et al., 2023)? These are widely adopted in world model literature and could be adapted to your setting as well. It would be worthwhile to include a comparison or at least a discussion to better show the difference between this approach and these alternatives to make it clear.  \n\n- A minor suggestion: I believe the reachability analysis should be included in the main paper, as it is crucial to justifying the objective function and its alignment with the theoretical guarantees.\n\n**W3**: [Empirical parts] Following FCDL, this paper evaluates the proposed method on environments including Robosuite and simulated chain and fork causal graphs. I believe these settings do demonstrate the effectiveness of the approach. However, from the perspective of the broader ambition expressed in the paper, I would encourage the authors to conduct more extensive evaluations to further strengthen the work.\n\n- The paper frequently highlights the potential of this approach for world models. However, there is no experimental validation in this direction. In world model settings, it is common to consider an observation function, whereas the current experiments seem to focus only on low-dimensional state inputs. While this is not a major issue, it is worth acknowledging and clarifying.\n\n- I believe it would be valuable to show the learned causal graphs in manipulation tasks with more complex structures, as done in CDL (see Fig. 22: https://proceedings.mlr.press/v162/wang22ae/wang22ae.pdf). This would help validate the method's ability to recover causal structures when more variables are involved.\n\n- I suggest including additional baselines that infer time-varying or context-specific causal structures directly from learned model components, such as transformer attention weights or the Jacobian matrices of MLP layers. These alternatives can often be more flexible within the world model paradigm, as they do not require learning a separate latent space. Comparisons with such baselines would provide stronger empirical support for the proposed approach. In particular, works like [Pitis et al., 2020], [Urpí et al., 2024], and potentially [Lei et al., 2024] (though I understand there is no code available for this one) appear relevant and could likely be tested with minimal additional effort.\n\n- Any empirical validation of the reachability of the intervention?\n\n**W4**: [Discussions on related works] First, to echo my earlier point, I believe FCDL and NCD deserve more explicit acknowledgement to make the contributions of this work clearer. Additionally, you mentioned that [Pitis et al., 2020] relies on strong prior knowledge. However, after reading that paper myself, I found the approach to be quite flexible. Could you clarify specifically what strong priors are required in their method?\n\n**W5**: [Minor] 1. The authors could consider including the corresponding graphical model and specifying additional assumptions on the graph structure (e.g., no instantaneous relationships) to make the framework more self-contained and easier to follow. 2. A typo in the expression for $\\mathcal{L}{\\text{quantization}}$? I suspect that in the second term $e_u$ should be ${E}_\\phi(X)$ instead? Correct me if I am mistaken.\n\n------\n\n**References**\n\nHwang, Inwoo, et al. \"Fine-grained causal dynamics learning with quantization for improving robustness in reinforcement learning.\" arXiv preprint arXiv:2406.03234 (2024).\n\nHwang, Inwoo, et al. \"On discovery of local independence over continuous variables via neural contextual decomposition.\" Conference on Causal Learning and Reasoning. PMLR, 2023.\n\nPathak, Deepak, et al. \"Curiosity-driven exploration by self-supervised prediction.\" International conference on machine learning. PMLR, 2017.\n\nKim, Kuno, et al. \"Active world model learning with progress curiosity.\" International conference on machine learning. PMLR, 2020.\n\nKauvar, Isaac, et al. \"Curious replay for model-based adaptation.\" arXiv preprint arXiv:2306.15934 (2023).\n\nPitis, Silviu, Elliot Creager, and Animesh Garg. \"Counterfactual data augmentation using locally factored dynamics.\" Advances in Neural Information Processing Systems 33 (2020): 3976-3990.\n\nLei, Anson, Bernhard Schölkopf, and Ingmar Posner. \"SPARTAN: A Sparse Transformer Learning Local Causation.\" arXiv preprint arXiv:2411.06890 (2024).\n\nUrpí, Núria Armengol, et al. \"Causal action influence aware counterfactual data augmentation.\" arXiv preprint arXiv:2405.18917 (2024).", "was_revised": true, "final_justification": "Rating Change: From 2 to 3.\n\n- Why the original rating of 2, and what are my current concerns?\n\nMy original score of 2 was primarily due to concerns about the technical novelty. Much of the method appears to be built directly upon FCDL, with added terms related to exploration (common in unsupervised RL and causal discovery works). This is not inherently problematic, but I felt the paper did not appropriately credit their base work. For instance, in Section 4.1, the way latent structures are learned is the same to FCDL (even for the notations), and in Section 4.3, the overall loss is almost identical. Section 4.2's entropy-based exploration is also a standard design in many discovery-oriented papers.\n\nA clearer and more transparent positioning, e.g., stating that the method builds on FCDL and introduces additional exploration terms, would be more appropriate. In its current form, this omission could lead to misunderstandings about the core technical contribution. I still have the same concerns after the rebuttal. \n\nIn my original review, I also raised other concerns, such as the choice of entropy computation, lack of clarity in some experiments, and the ambiguous use of the term \"world model.\"\n\n2. How did the rebuttal address these issues?\n\nThe authors made a genuine effort to address most of the concerns, and I truly appreciate that. This is the main reason I’ve raised my rating. While there are still points of disagreement, it’s clear the authors are working toward improving the paper.\n\nAside from the FCDL-related concerns and the framing of the world model, I believe the other points I raised have been sufficiently addressed.\n\n3. Why not a higher score?\n\nDespite the progress, I still have reservations about several aspects:\n\n- The current draft heavily borrows from FCDL without clearly acknowledging its influence. The modeling components in Section 4 require a thorough rewrite to accurately attribute prior work and clarify the incremental contribution. See the above points I mentioned. \n\n- The use of the term \"open-ended tasks\"in the rebuttal, such as open-ended games, multi-agent systems, or LLM-based dialogue, feels overstated. While this serves as an inspiring motivation, the method does not demonstrate any empirical advantages in these domains, and I suggest removing or reframing this claim.\n\n- Regarding the \"world model\" framing: I believe the paper overclaims its relevance. First, treating actions entirely as interventions is questionable. Second, there is no clear evidence that the learned model is useful for decision-making. While that may be a long-term goal, more technical depth is needed to support positioning this as a world model paper. These issues prevent me from justifying a higher rating or acceptance at this point.\n\n4. Recommendation and suggestions:\n\nGiven these concerns, I remain on the fence about it. I understand that other reviewers are more positive, so I’m happy to defer to the majority opinion and engage in further discussion. If the consensus leans toward acceptance, I strongly recommend requiring the authors to revise the framing and claims, especially those regarding world models and open-ended tasks, and importantly, to significantly rewrite Section 4 to clarify the model's contributions and properly acknowledge the influence of FCDL.", "rebuttal": "We thank the reviewer for acknowledging the importance of meta-causal graph (MCG), the effectiveness of our methodology, and the strength of our theoretical results on the identifiability of MCG. Below are our responses. Issues you specifically raised are marked with $\\star$.\n\n> $\\star$ **W1: Simlarity with FCDL/NCD**\n\nFCDL and NCD are good works, but our method is totally different. \n\n**Different Motivation**：\n- FCDL & NCD are **observation-based methods** restricting themselves to **passive learning** and limits their capacity to truly explore or adapt in an open-ended world, while our method is **intervention-based**, which helps us to deal with the open-ended world problem.\n- Our goal is to **actively discover** the unseen meta state in open-ended worlds to learn a **MCG** (line 39-40).\n\n**Different Formulation and theoretical contributions:**\n- **FCDL & NCD do not model the minimal causal graph and cannot guarantee identifiability**; their approach can lead to non-identifiable solutions due to infinite subgraph combinations. While we develop a theory framework to ensure identifiability of MCG.\n- FCDL/NCD treats actions as part of **observational data**, which does not allow for active intervention. Our method treats actions as **interventions**, enabling true active intervention for causal discovery in open worlds.\n\n> $\\star$ **W2.a: Tractability of entropy computation and robustness to inaccurate posterior dist**\n\n**The computation of entropy in Eq. 2 is tractable.**\n- Entropy is calculated as the sum of the entropies of Bernoulli variables, where each variable $M[i, j]$ represents the existence of an edge $i\\to j$ in the causal subgraph and the corresponding Bernoulli probabilities $\\hat{M}[i, j]=P(M[i, j]=1)$ are predicted by the neural network. \n- Implementation details are provided in code file `fcdl/utils/utils.py`, function `calculate_entropy`.\n\nWe **added experiments testing robustness to posterior inaccuracies** by introducing disturbances. These results show the robustness of our method.\n|Acc↑|n=2|n=4|n=6|\n|-|-|-|-|\n|FCDL|57.8|49.3|47.7|\n|**MCG (disturbed)**|**60.1**|**51.0**|**52.7**|\n|**MCG**|**63.2**|**50.5**|**50.0**|\n> $\\star$ **W2.b: Discussion with other curiosity methods**\n\nGenerally, curiosty-based RL methods require designing reward encouraging exploration, such as [Pathak et al., 2017], [Kim et al., 2020] and [Kauvar et al., 2023]. Compared with these methods, our curiosity-driven intervention is tailored for discovering new meta states in the MCG, which is unique to our framework, while traditional methods do not target unseen meta states.\n\nWe will discuss it more in the revised version.\n> **W2.c: Move the reachability analysis from appendix to the main paper**\n\nThank you for the suggestion. We will move the reachability analysis from appendix to the main paper in the revised version.\n\n> **W3.a: Compared with traditional world models, why not use an observation function to obtain representations?**\n\n“World model” is a broad term. Not all world models operate directly from high-dimensional observations; some, including [e.g., Liu et al.], focus on **extracting certain semantic representations** from the state space. Similarly, we **learn unknown meta state representations** and their corresponding causal graphs directly from the state space.\n> $\\star$ **W3.b: Display the learned causal graphs in more complex tasks**\n\nWe have **added a manipulation task with a more complex structure**. Due to space constraints, we present only a subset of the graph. This experiment validates our method’s ability to recover causal structures in settings with 15 variables. Larger one will be added in next version.\n\n**MCG**\n||obj 0|obj 1|obj 2|obj 3|obj 4|\n|-|-|-|-|-|-|\n|obj 0|1.0|0.0|0.0|0.0|0.0|\n|obj 11|1.0|1.0|0.5|0.5|0.6|\n|obj 12|0.6|0.7|0.9|0.0|0.1|\n|obj 13|1.0|0.9|0.7|0.0|0.4|\n\n**Ground truth**\n\n||obj 0|obj 1|obj 2|obj 3|obj 4|\n|-|-|-|-|-|-|\n|obj 10|1.0|0.0|0.0|0.0|0.0|\n|obj 11|1.0|1.0|1.0 |1.0|1.0|\n|obj 12|1.0|1.0|1.0|0.0|0.0|\n|obj 13|1.0|1.0|1.0|0.0|0.0|\n\n> $\\star$ **W3.c: Add additional baselines**\n\nWe **added baselines you suggested**: causal discovery using transformer attention weights (\"transformer\") and the Jacobian matrices of MLP layers (\"sandy-mixture\") from [Pitis et al., 2020] and [Urpí et al., 2024], for comparison. The results are shown below. Our method outperforms new baselines, demonstrating that our active intervention strategy improves the accuracy of the recovered causal graph.\n\n|Acc↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|transformer|24.45|22.49|21.41|29.62|30.37|29.78|\n|sandy-mixure|31.93|32.47|33.72|30.08|29.31|27.43|\n|**MCG**|**63.18**|**50.47**|**50.04**|**51.99**|**49.78**|**49.69**|\n\n|Acc↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|transformer|6.20|6.13|6.71|6.67|6.73|7.40|\n|sandy-mixure|6.48|6.10|6.50|7.03|6.73|7.17|\n|**MCG**|**14.65**|**14.06**|**13.28**|**13.82**|**12.49**|**12.45**|\n> $\\star$ **W3.d: Empirical validation of reachability**\n\nWe **added experiments comparing causal graphs under limited vs. full reachability**. Limited reachability yields less accurate graphs (higher SHD, lower F1).\n||SHD↓|F1↑|\n|-|-|-|\n|MCG (limited reachability)|17.31|0.77|\n|MCG|11.69|0.84|\n\n> **W4.a: Discussions on related works**\n\nWe acknowledge the valuable contributions and inspiration of FCDL/NCD in causal discovery and will add more discussion in the revision. We believe that key conceptual differences have already been clarified in our response to **Comment W1**. We provide other comparisons here:\n- Observation-based methods like FCDL and NCD are limited to Markov equivalence classes, while intervention-based methods (Ours) can **actively reduce the Markov equivalence class and uniquely identify causal relations** that are otherwise unidentifiable.\n- We introduce **$L_{mask}$ loss** (line 197-198) that directly updates causal structure based on intervention outcomes while FCDL & NCD have no mechanism to verify or update causal beliefs through interventions.\n\n> **W4.b: Citation for [Pitis et al., 2020]**\n\nApologies for the citation error. The intended reference was another paper, [Pitis et al., 2022], which discusses methods that require domain knowledge or expert trajectories. We will correct this mistake in the revised version.\n\n> **W5: Typos and clarification of theoretical aspects of the paper**\n\nThank you for your careful reading. We will improve the clarity of the theoretical aspects of the paper and fix the typos.\n\n> **Q1.a: How to integrate and learn context-specific causal discovery approach in world model efficiently**\n- Our approach uses a curiosity-driven exploration strategy to **actively discover unknown meta states**. Specifically, the agent **maps each state to a learned representation** that jointly captures the meta state and its corresponding causal graph.\n- Based on the discovered causal graphs, we learn a transition model conditioned on meta states. Here, the **MCG** serves as **structured world knowledge**, while the **learned transition model functions** as the **world model** (see Definition 3 and line 204-205).\n- We added experiments to show efficiency of our method. We compare other methods (GNN and MLP) with longer training rounds and more parameters. The result is shown below:\n\n**Comparison with _Same_ Model Size**\n\n||Acc n=2|Acc n=4|Acc n=6|Reward n=2|Reward n=4|Reward n=6|\n|-|-|-|-|-|-|-|\n|GNN (30k round)|38.13|25.81|20.75|6.83|6.06|6.83|\n|GNN (25k round)|36.37|26.55|20.73|6.84|5.98|6.19|\n|MLP (30k round)|30.47|28.94|27.84|6.20|6.13|6.71|\n|MLP (25k round)|31.47|30.32|30.42|6.18|6.03|6.23|\n|**MCG** (10k round)|**47.01**|**48.37**|**46.35**|**11.96**|**10.83**|**10.70**|\n|**MCG** (12k round)|**48.30**|**48.40**|**47.28**|**12.45**|**10.77**|**10.84**|\n|**MCG** (15k round)|**71.56**|**59.17**|**49.58**|**14.65**|**14.06**|**13.28**|\n\n**Comparison with _Same_ Round**\n||Acc n=2|Acc n=4|Acc n=6|Reward n=2|Reward n=4|Reward n=6|\n|-|-|-|-|-|-|-|\n|MLP (parameter $\\times$ 3.71)|24.65|25.05|23.71|6.29|6.50|7.93|\n|GNN (parameter $\\times$ 1.66)|37.58|30.61|30.05|6.34|6.49|7.90|\n|**MCG** (parameter $\\times$ 1)|**71.56**|**59.17**|**49.58**|**14.65**|**14.06**|**13.28**|\n\nOur method achieves strong performance with less training data and smaller models, highlighting its efficiency.\n> **Q1.b: How should actions be incorporated into the framework?**\n\nIn our framework, actions are treated as **interventions** for causal discovery. \n\n> **Q1.c: Additional structural assumptions needed?**\n\nWhile not all actions can intervene on every state (i.e., not all states are reachable by intervention), we explicitly make a **reachability assumption** to ensure identifiability of the underlying structure.\n\n> **Q1.d: Can you leverage this structure discovery as part of reward learning?**\n\nOur curiosity-driven reward (Eq. 2) is explicitly designed to encourage MCG structure discovery of meta-states $u=C(x)$ and subgraphs $\\mathcal{G}_u$.\n\n> **Q2.a: Clarification of observation equivalence**\n\nThank you for your comment. We believe there may be a slight misunderstanding. Observational equivalence refers to the condition where two mappings from states to meta states produce the same partitioning of observations, possibly up to relabeling. We will clarify this in the revised version.\n\n> **Q2.b: How can imperfect representation affect the identifiability of meta-states and the causal graph?**\n\nThank you for your comment. We agree that learned latents can be imperfect. Empirically and theoretically, our study on (Fig. 2 and Thm. 2) shows over‑parameterisation does not hurt performance much while under-parameterisation will lead to degradation of performance. We will add more discussions on this point in the revised version.\n\nYang Liu, et al. Learning world models with identifiable factorization. NeurIPS, 2023.\n\nSilviu Pitis, et al. Mocoda: Model-based counterfactual data augmentation. NeurIPS, 2022."}, {"review_id": "3KjXwctM48", "reviewer": "Reviewer_6jMv", "summary": "The paper introduces a theoretically grounded, curiosity-driven agent that actively discovers a set of context-dependent causal graphs, which are called meta-causal graphs, leading to more robust world models and improved downstream performance over existing passive or single-graph approaches.\n\nThe agent observes state-trajectories and occasionally intervenes on individual variables through a curiosity (entropy reduction) mechanism. From this mix of passive observations and active single-variable interventions it learns by (i) clustering observations into a higher level meta-state, (ii) estimate the causal skeleton and edge directions of observation variables given the inferred meta-state, and (iii) fit a transition predictor whose factorisation follows the learned graph (which can be used for model-based planning; cf world model).\n\nThe paper backs the method with theoretical results that show the agent can, in principle, identify the structure of the environment: Meta-state identifiability (Theorem 1), and Causal-edge identifiability via interventions (Proposition 1). This is a priori not possible by methods which utilize a single causal graph, i.e., which are insensitive to the context variable for the environment. \n\nThe paper benchmarks the learned world models on two environments: a Chemical environment and a more challenging Magnetic robot-arm task. The authors benchmark out-of-distribution prediction accuracy on the chemical environment and cumulative reward in both environments, showing that the proposed method yields the highest performance on all metrics and against several competing baselines.", "questions": "Suggestion: Make the paper more accessible to the reinforcement learning literature that does not involve causality. Taking the example of Figure 1, for instance, the paper explains this with the causal graph push -> open door or push -> door remains closed. Therefore, it makes the point that we need a contextual meta-state that tells us which of these graphs is applicable in any given situation. A standard model-based reinforcement learning model of this would be a POMDP where state of door at time t (which could be open or closed) transitions to state of door at time t+1 (which could be open or closed) based on the action push. This (potentially stochastic) transition would depend on an extra state (the meta-state in the author's words; door locked or unlocked), at time t, that determines the consequence of push at time t+1. In other words, it seems to me that we can recapitulate this simple example in a standard POMDP, where one of the states is a contextual state. Perhaps one advantage of this formulation is that we can have fuzzy causal arrows in the sense that the push action could have a stochastic consequence on the state of the door. (Not that this is needed here, but it's useful in other types of situations). It seems to me that the causal formulism is handy because it produces these nice interpretable diagrams that do not have to invoke stochastic processes like POMDPs. The reader from model-based reinforcement learning will certainly wonder at the differences/respective features between these two frameworks. One classic paper, for instance, shows the equivalence between causal diagrams and interventions with probabilistic graphical models and intervention nodes (like POMDPs). \n\nCausality with Gates, Winn, John, Proceedings of the Fifteenth International Conference on Artificial Intelligence and Statistics, 2012.\n\nI do not mean to criticize the causal framework used in the paper, just to note that the exposition may be unnecessarily divorced from standard model-based reinforcement learning, which limits understanding/clarity by a broader audience and potential impact. Perhaps one path of least action would be to clarify these discrepancies/equivalences in the related works section or somewhere else. \n\nIncrease emphasis of unique contribution: a clear contextualization within the model-based reinforcement learning literature that does not use the causal framework and a mention of related work that learns latent variable dynamic models (such as but not limited to POMDPs) through curiosity-driven actions would help contextualize and emphasize the unique contributions of this paper in the broader RL Literature.\n\nSome words on efficiency and scale: one strength of the causal approach is parameter and compute efficiency (promoted by your sparsity regularizer and curiosity-driven exploration). Unless I missed it, there is no discussion of model efficiency in the paper. Is there anything that can be said here? Same goes with scalability. The results on the environments tested are impressive. So it begs the question: what are the limitations to scaling to more complex environments? When would you expect the method to scale and when would you expect it not to scale? This would be very helpful for future work.", "limitations": "Yes.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Clarity: I found the abstract not very clear. Generally, I find the paper clear to a causality audience, but less so for a reinforcement learning audience. I added some suggestions to improve this in the next section. \n\nQuality: the paper combines experimentation with nice theoretical results regarding identifiability and clearly acknowledges some of the limitations of the proposed approach. The curiosity-driven intervention approach is a refreshing addition to the classical causal framework. Curiosity in causal learning is not a new idea and indeed early works from the early 2000s considered this, e.g. [1]. However, it seems to have fallen out of fashion, and most works in causal learning do not consider uncertainty (only maximum likelihood inference) and therefore cannot leverage uncertainty for interventions. The proposed implementation of curiosity for uncertainty reduction is novel to me. \n\n[1] Active Learning of Causal Bayes Net Structure by Murphy (2001).\n\nOriginality: the originality is hard to fully assess. Certainly many things are original, but a clear contextualization within the model-based reinforcement learning literature that does not use the causal framework and a mention of related work that learns latent variable dynamic models (such as but not limited to POMDPs) through curiosity-driven actions would help contextualizing and emphasizing the unique contributions of this paper to the Broad RL Literature.\n\nSignificance: the experiments and theoretical results are compelling. The method achieves the best out of distribution prediction accuracy in the first experiment, and significantly boosts downstream episode reward in both experiment, beating all competing methods across all metrics considered. The experiments are rigorous: Ablations show that removing either the curiosity-based reward or the causal mask loss hinders both metrics, confirming their necessity, and comparison is done with several baselines. Furthermore, theoretical results such as Theorem 2, which predicts that overparameterisation never hurts (even if the agent allocates more codebook slots than there are real contexts, every true context still gets its own slot and extra slots fold onto them, leaving the recovered sub-graphs equivalent to ground truth), is validated in practice. One criticism it is hard to fully determine the significance of the paper since there are no discussion on efficiency, run-time and potential feasibility to scale to more complex tasks. In particular, it is unclear whether the method would really apply to real-world problems.", "was_revised": true, "final_justification": "I thank the authors for providing detailed responses to my comments. As it stands, and with the new experiments, the paper seems strong, and I am happy to recommend acceptance. Indeed, their method almost consistently outperforms alternatives, and the authors included an experiment where their method works in a partially observable domain, edging closer to real-world applicability. My main point of concern involving the clarity of the paper and its contextualization within the model-based RL literature has been promised to be addressed by the authors in their camera ready version. However, I do not know the details of these edits and therefore cannot increase my clarity score at this stage. All in all, I recommend acceptance of the paper.", "rebuttal": "We thank the reviewer for highlighting that our curiosity-driven agent and meta-causal graph framework provide a theoretically grounded and practically effective approach to robust world modeling, with compelling empirical and theoretical results that advance causal learning beyond traditional single-graph or passive methods. Below are our responses to your comments and suggestions. \n\n> **Discussion differences between MBRL under POMDP settings and intervention-based causal method  [Winn et al.]**\n\nWe appreciate the reviewer’s thoughtful suggestion.\n\nOur framework can be viewed as a special case of MBRL, but more detailed differences are listed below:\n- **Unified World Representation**: Our framework maintains a single meta-causal graph per environment, **avoiding redundant definitions** that arise from time-series changes in POMDPs.\n- **Explicit Causal Structure**: Meta-causal graphs directly **represent variable-level causal relationships** in each context, whereas POMDPs only model transition probabilities without explicit causal structure.\n- **Identifiability Guarantees**: The meta-causal graph framework provides **theoretical guarantees for identifying both meta-states and causal structures, ensuring reliable causal discovery**. In contrast, POMDPs generally lack such identifiability guarantees for their latent structure.\n\nCompared to the intervention-based causal method mentioned by the reviewer (Winn et al.), their approach is not RL-based and does not incorporate curiosity-driven mechanisms to encourage intervention-based exploration, as our method does.\nA detailed discussion of the similarities and differences with more related works will be included in the revised version.\n\n> **Efficiency and scalability of our method**\n\nTo demonstrate the efficiency and scalability of our method, we **added new experiments** comparing it to non-causal baselines (MLP and GNN) across **different model sizes and training rounds**. The metrics are next-state prediction accuracy and the average episode reward obtained by model-based planning on downstream tasks, evaluated under varying levels ($n$) of noise (lines 265–268).\n\n**Comparison with _Same_ Model Size**\n\n||Acc n=2↑|Acc n=4↑|Acc n=6↑|Reward n=2↑|Reward n=4↑|Reward n=6↑|\n|-|-|-|-|-|-|-|\n|GNN (30k round)|38.13|25.81|20.75|6.83|6.06|6.83|\n|GNN (25k round)|36.37|26.55|20.73|6.84|5.98|6.19|\n|MLP (30k round)|30.47|28.94|27.84|6.20|6.13|6.71|\n|MLP (25k round)|31.47|30.32|30.42|6.18|6.03|6.23|\n|**MCG** (10k round)|**47.01**|**48.37**|**46.35**|**11.96**|**10.83**|**10.70**|\n|**MCG** (12k round)|**48.30**|**48.40**|**47.28**|**12.45**|**10.77**|**10.84**|\n|**MCG** (15k round)|**71.56**|**59.17**|**49.58**|**14.65**|**14.06**|**13.28**|\n\n**Comparison with _Same_ Round**\n||Acc n=2↑|Acc n=4↑|Acc n=6↑|Reward n=2↑|Reward n=4↑|Reward n=6↑|\n|-|-|-|-|-|-|-|\n|MLP (parameter $\\times$ 3.71)|24.65|25.05|23.71|6.29|6.50|7.93|\n|GNN (parameter $\\times$ 1.66)|37.58|30.61|30.05|6.34|6.49|7.90|\n|**MCG** (parameter $\\times$ 1)|**71.56**|**59.17**|**49.58**|**14.65**|**14.06**|**13.28**|\n\nThe results show that our method can achieve comparable performance with **fewer training data and smaller model sizes**, demonstrating its efficiency and scalability.\n\n> **Scalability potential and challenges of our method**\n\n- **Our theoretical framework is general and can, in principle, scale to more complex tasks**, including those with delayed or partial observations.\n    - We **added a new experiment** to demonstrate that our method can be naturally extended to open-world settings where **variables are partially observable**.\n    - We simulate **the appearance of new variables** by initially masking a subset of variables and gradually removing these masks during training, so that the set of observed variables increases over time.\n    - The metric is next-state prediction accuracy, evaluated under varying levels ($n$) of noise (lines 265–268).\n|Accuracy↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|FCDL|60.10|52.97|45.80|46.37|**47.88**|46.62|\n|MLP|31.56|31.67|31.93|25.95|24.71|21.66|\n|**MCG**|**66.94**|**59.19**|**46.35**|**49.93**|45.41|**48.09**|\n\n- **Engineering challenges such as model architecture design and efficient exploration strategies** may require further adaptation as environments become larger or more complex. Addressing these challenges will be an important direction for future work.\n    - We **added a manipulation task with a more complex structure**. Due to space constraints, we present only a subset of the graph. This experiment validates our method’s ability to recover causal structures in settings with 15 variables. We will extend this experiment to larger environments in next version.\n\n        **Causal graph recovered by MCG**\n        ||obj 0|obj 1|obj 2|obj 3|obj 4|\n        |-|-|-|-|-|-|\n        |obj 0|1.0|0.0|0.0|0.0|0.0|\n        |obj 11|1.0|1.0|0.5|0.5|0.6|\n        |obj 12|0.6|0.7|0.9|0.0|0.1|\n        |obj 13|1.0|0.9|0.7|0.0|0.4|\n\n        || obj 4 |obj 5|obj 6|obj 7|obj 8|\n        |---|---|---|---|---|---|\n        |obj 4| 1.0 | 0.0 | 0.0 | 0.0 | 0.0 |\n        |obj 5| 0.1 | 1.0 | 0.0 | 0.0 | 0.0 |\n        |obj 6| 0.2 | 0.1 | 1.0 | 0.0 | 0.0 |\n        |obj 7| 0.3 | 0.0 | 0.2 | 1.0 | 0.0 |\n        |obj 8| 0.3 | 0.1 | 0.5 | 0.5 | 1.0 |\n\n        **Ground truth causal graph**\n\n\n        ||obj 0|obj 1|obj 2|obj 3|obj 4|\n        |-|-|-|-|-|-|\n        |obj 10|1.0|0.0|0.0|0.0|0.0|\n        |obj 11|1.0|1.0|1.0 |1.0|1.0|\n        |obj 12|1.0|1.0|1.0|0.0|0.0|\n        |obj 13|1.0|1.0|1.0|0.0|0.0|\n\n        ||obj 4|obj 5|obj 6|obj 7|obj 8|\n        |---|---|---|---|---|---|\n        |obj 4| 1.0 | 0.0 | 0.0 | 0.0 | 0.0 |\n        |obj 5| 0.0 | 1.0 | 0.0 | 0.0 | 0.0 |\n        |obj 6| 0.0 | 0.0 | 1.0 | 0.0 | 0.0 |\n        |obj 7| 0.0 | 0.0 | 0.0 | 1.0 | 0.0 |\n        |obj 8| 0.0 | 0.0 | 0.0 | 0.0 | 1.0 |\n\n> **Increase emphasis of unique contribution compared with MBRL**\n\nThank you for the valuable suggestion.\nOur framework can be viewed as a special case of MBRL, but more detailed differences are listed below:\n- We agree that highlighting the importance of **causal understanding** distinguishes our approach from standard model-based RL methods, which often struggle to scale and to continually adapt to new causal structures in open-ended environments.\n- Our method operates within an MBRL framework, but the meta-causal graph provides a **structured, causal understanding of the world**—enabling the agent to develop a **causally grounded MDP that adapts to changing dynamics**.\n- We will include a more detailed discussion of connections to both standard MBRL and curiosity-driven latent variable models (e.g., POMDPs) in the revised version."}, {"review_id": "wdC03dXzZb", "reviewer": "Reviewer_n7bY", "summary": "This paper focuses on the representation and identification of a world model in a causal lens. \n\nIn the first part, it introduces the key concept -- Meta-Causal Graph. It is a collection of causal subgraphs, and each subgraph corresponds to a distinct meta status. The identifiability of both meta status and Causal Subgraphs is presented in detail.\n- The identifiability of meta status relies on the existence of the (state)-to-(meta state) mapping: that is, one can distinguish observations from different meta-subgraphs by observing their values.\n- The identifiability of subgraphs relies on sufficient intervention targets applied on at least one directed edge.\n\nIn the second part, it develops a Causality-Seeking Agent to (1) exploit existing experience to estimate the Meta-Causal Graph, including the status mapping and each subgraph; (2) actively explore the environment by taking actions. The key design here is a curiosity-driven reward. It is defined by the Shannon entropy of the learnable distributions on edges. This makes sense because higher entropy implies higher uncertainty.\n\nThe paper combines the two parts together: the agent decides what actions (or interventions) to take based on the current estimated Meta-Causal Graph; and the Meta-Causal Graph would be updated based on the collected data from actions. The Meta-Causal Graph serves as a world model for the agent.\n\nIn evaluation, it uses the Chemical environment setting. (1) OOD states in the Chemical environment; (2) rewards on the downstream task. In this part, while the results seem promising, some details are missing, for example, what are the downstream tasks in the Chemical environment? What does it mean to \" predict transition dynamics for model-based reinforcement learning\"? I checked the appendix. I cannot open the link to the anonymous code repo.", "questions": "1. In Definition 3 and Assumption 1, it assumes the meta state $u$ can be determined from observations $\\mathbf{x}$. Wouldn't this be a very strong assumption? \n\n2. It seems not to include the usual assumptions for causal discovery, like faithfulness and Markovness conditions. Am I missing something?\n\n3. While the experiment provides a careful evaluation of downstream tasks, it might be necessary to evaluate the causal discovery metrics (SHD, SID, F1 score, etc) compared with other baselines (methods that learn multiple causal graphs from data).\n\n4. In reality, conducting interventions can be expensive. How would the performance be influenced if: (1) the round of Algorithm 1 is limited, (2) the size of samples from the intervened distribution is small?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n- The paper proposes a Meta-Causal Graph to serve as a world formulation. This setting considers environmental conditions shift by introducing the (status)-(meta status) mapping component.\n- Section 3 presents identifiability on both meta status and Causal Subgraphs.\n\n\n**Weaknesses**\n- I have concerns that the proposed Meta-Causal Graph may still have limitations on serving as a world model. (1) The assumption on the existence of such (status)-(meta status) mapping needs to be further justified. (2) The model assumes a fixed set of variables, which may not hold in an open world: one may encounter new variables/objects during exploration.\n- The details in the experimental setting are not sufficient. For example, what are the downstream tasks in the Chemical environment? What does it mean to \" predict transition dynamics for model-based reinforcement learning\"? Also,  it might be necessary to evaluate the causal discovery metrics (SHD, SID, F1 score, etc) compared with other baselines (methods that learn multiple causal graphs from data).\n\n\n**Explanation on the decision**:\n\nThis paper contains sufficient technical details, like the identifiability of each component in the Meta-Causal Graph model, the parameterization, the reward model, and the loss function in the agent part.\n\nThe paper presents insight on how to combine a world model (based on a collection of causal graphs ) and an agent that can take actions to collect new data from actively applying interventions.\n\nLimitation includes insufficient details in the experiment section, and the assumptions on the existence of (status)-(meta status) mapping, and a fixed set of variables across all environments. \n\nI recommend a Borderline acceptance if the experiment details can be clarified, and the mentioned limitation can be fairly discussed in the rebuttal phase. Otherwise, it is possible to decrease the score.", "was_revised": true, "final_justification": "The original concerns: insufficient details in the experiment section, and the assumptions on the existence of (status)-(meta status) mapping, and a fixed set of variables across all environments. \n\nBased on the current response, the issues were resolved:\n- a fixed set of variables across all environments: the author provides additional results to demonstrate the potential to handle new variables during explorations.\n- (partially solved) The author includes additional metrics to augment the results further.\n\nIssues remain unresolved: \n- (major) the existence of (status)-(meta status) mapping: The authors said it would *necessarily exist when all variables are observable*. I provide a counter-example, and also provide further suggestions.\n- (major) the theoretical assumptions. The authors say that the stated assumptions are sufficient, while not providing concrete explanations. I provide a follow-up comment to further explain my concern with evidence from existing literature.\n- (minor) The author hasn't updated the details behind table 2.\n\n--- \nweight assigned: \n- one major issue not solved: -1.\n- one major issue partially solved: -0.5.\n- one minor issue not solved: -0.5.\n\nTherefore, the updated score is: 4 (initial) - 1 - 0.5 - 0.5 = 2\n\n\n-----\n\n\n**After the second response**\n- The issues are carefully addressed/clarified, --> 4", "rebuttal": "We thank the reviewer n7bY for highlighting the strengths of our approach, especially the introduction of the Meta-Causal Graph for serving as the world model and our detailed analysis of identifiability for both meta states and causal subgraphs. Below are our responses to your comments and suggestions. \n\n> **W1&Q1: Existence of state to meta state mapping when new variables appear during exploration in open-ended world. Would the observable assumption too strong?**\n\n- Firstly, we would like to highlight that **the mapping from observation $x$ to meta state $C(x) = u$ necessarily exists when all variables are observable.**\n- This is because each state $x$ is governed by a unique underlying causal subgraph $\\mathcal{G}_u$, and each meta state $u$ indexes a distinct subgraph.\n- Secondly, **our method can be naturally extended to open-world settings where new variables or objects are encountered during exploration**.\n- To support this, we **added additional experiments (detailed setup and results below) showing performance across time steps as new variables are revealed**. The results show that our method outperforms baselines, highlighting its strong exploratory and adaptive capabilities.\n\n**Setup**\n\n- We simulate **the appearance of new variables** by initially masking a subset of variables and gradually removing these masks during training, so that the set of observed variables increases over time.\n- The metric is next-state prediction accuracy, evaluated under varying levels ($n$) of noise (lines 265–268).\n\n**Result**\n\n|Accuracy↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|FCDL|60.10|52.97|45.80|46.37|**47.88**|46.62|\n|MLP|31.56|31.67|31.93|25.95|24.71|21.66|\n|**MCG**|**66.94**|**59.19**|**46.35**|**49.93**|45.41|**48.09**|\n\n> **W2.a: What is the downstream task in the Chemical environment in paper**\n\nIn the Chemical environment downstream task:\n- The agent starts from a random color configuration and must **transform a 10-node graph to match a target color pattern.**\n- **Actions intervene on nodes, changing their color and that of all their descendants** according to the hidden causal graph.\n- The episode reward is the negative **Hamming distance to the target (0 for a perfect match)**.\n\n> **W2.b: The meaning of \"predict transition dynamics for model-based reinforcement learning\"**\n\n\"Predict transition dynamics for model-based reinforcement learning\" means **predict future states with the learned model for decision making**.\n\nWe will clarify these in the revised version.\n\n> **Q1：The observable assumption too strong for get meta-state u**\n\nWe detailedly discuss it in W1&Q1.\n\n> **Q2: Clarification of assumptions to support causal discovery**\n- Our theoretical results are established under our stated assumptions, which are **sufficient** for our proofs (as also noted by Reviewer hzVP).\n- We appreciate your suggestion and will consider the extension of our theory to more general conditions, including faithfulness and Markov assumptions, in next version.\n\n> **Q3: Evaluate the causal discovery metrics (SHD, SID, F1 score, etc)**\n\nWe add experimental results for these metrics. The results are shown below:\n||SHD↓ (Fork)|SID↓ (Fork)|F1↑ (Fork)|SHD↓ (Chain)|SID↓ (Chain)|F1↑ (Chain)|\n|-|-|-|-|-|-|-|\n|NCD|21.67|21.67|0.77|19.15|19.15|0.76|\n|FCDL|17.05|17.05|0.77|19.39|19.39|0.76|\n|**MCG**|**16.59**|**16.59**|**0.78**|**11.64**|**11.64**|**0.84**|\n\n> **Q4: Intervention efficiency analysis when (1) the round of Algorithm 1 is limited, (2) the size of samples from the intervened distribution is small?**\n\nWe **added an experiment to evaluate the intervention efficiency by reducing the number of interventions**. \n- In this experimental setting, **data is collected through interventions at each round**; therefore, **reducing the number of rounds** corresponds to **a direct reduction in available intervention data**. \n- The metric is the average episode reward from model-based planning on downstream tasks, evaluated under varying levels ($n$) of noise (lines 265–268). The results are shown below:\n\n|Reward↑|Fork (n=2)|Fork (n=4)|Fork (n=6)|Chain (n=2)|Chain (n=4)|Chain (n=6)|\n|-|-|-|-|-|-|-|\n|FCDL (100% data, 15k rounds)|13.98|13.36|12.91|11.16|10.39|10.62|\n|NCD (100% data, 15k rounds)|10.95|9.11|9.11|9.60|8.86|10.32|\n|**MCG (60% data, 10k rounds)**|11.96|10.83|10.70|11.97|11.05|11.38|\n|**MCG (80% data, 12k rounds)**|12.45|10.77|10.84|12.41|10.87|11.00|\n|**MCG (100% data, 15k rounds)**|**14.65**|**14.06**|**13.28**|**13.82**|**12.49**|**12.45**|\n\nThe results show that our method can achieve comparable performance with **fewer interventions, demonstrating its efficiency in causal discovery**.\n\n> **Accessibility of the anonymous code repo**\n\nWe have checked the anonymous code repository and confirmed that it is accessible. We kindly suggest trying again at a later time."}]

Author rebuttal:

Meta-review: The paper proposes a method for learning a world model that adapts to changing causal dynamics. An agent actively explores using a curiosity-driven mechanism to find different situations (meta-states) and learns a separate graph of cause-and-effect relationships for each one. Reviewers acknowledge several strengths including theoretical contributions on identifiability, the importance of the problem setting, and the novelty of the curiosity-driven intervention approach. The weaknesses noted include: (1) attribution concerns regarding insufficient acknowledgment of FCDL, (2) insufficient experimental details and scope, (3) clarity issues for broader audiences, and (4) overstated claims about open-ended tasks and world model framing without sufficient empirical validation.

In the rebuttal, the authors provided new experiments, comparisons, and clarifications and acknowledged limitations regarding overstated claims. While three of four reviewers gave positive recommendations, key concerns remained unresolved regarding attribution, framing, and experimental scope limitations.

The Area Chair agrees with the majority of reviewers and recommends acceptance. Authors are advised to revise based on reviewers' comments, particularly by rewriting Section 4 to clearly distinguish novel contributions from prior work, moderating claims about world models and open-ended applications to match experimental scope, and improving accessibility for the broader ML community.
Decision: Accept (poster)