Language Agents Meet Causality -- Bridging LLMs and Causal World Models
Under review as a conference paper at ICLR 2025

LANGUAGE AGENTS MEET CAUSALITY – BRIDGING
LLMS AND CAUSAL WORLD MODELS

Anonymous authors
Paper under double-blind review

ABSTRACT

Large Language Models (LLMs) have recently shown great promise in planning and
reasoning applications. These tasks demand robust systems, which arguably require
a causal understanding of the environment. While LLMs can acquire and reflect
common sense causal knowledge from their pretraining data, this information is
often incomplete, incorrect, or inapplicable to a specific environment. In contrast,
causal representation learning (CRL) focuses on identifying the underlying causal
structure within a given environment. We propose a framework that integrates CRLs
with LLMs to enable causally-aware reasoning and planning. This framework
learns a causal world model, with causal variables linked to natural language
expressions. This mapping provides LLMs with a flexible interface to process and
generate descriptions of actions and states in text form. Effectively, the causal world
model acts as a simulator that the LLM can query and interact with. We evaluate
the framework on causal inference and planning tasks across temporal scales and
environmental complexities. Our experiments demonstrate the effectiveness of the
approach, with the causally-aware method outperforming LLM-based reasoners,
especially for longer planning horizons.

1

INTRODUCTION

Large Language Models (LLMs) have emerged as powerful tools for a wide range of tasks, from
natural language understanding to complex problem-solving (Brown et al., 2020; Radford et al., 2019;
Liu et al., 2023b). Recent work has explored the use of LLMs as action agents for planning and
reasoning tasks, showing promising results in improving task-specific, downstream performance (Ahn
et al., 2022; Hao et al., 2023; Huang et al., 2023). These approaches primarily rely on the model’s
ability to extract common-sense causal information stated in its training data (Zeˇcevi´c et al., 2023).
While LLMs can reflect general beliefs and correlations, this information may be incomplete, incorrect,
or inapplicable in specific environments. This poses challenges for LLMs in novel or complex
situations, particularly in dynamic environments where accurate modeling of action consequences is
crucial (Valmeekam et al., 2023; Kambhampati et al., 2024).

Causal representation learning (CRL) aims to identify the underlying causal structure of
data (Schölkopf et al., 2021). By separating and identifying latent causal factors, CRL enables
models to reason about the effects of interventions and counterfactuals. Recent theoretical work
provides justification for causal representation learning, showing it is necessary for achieving strong
robustness guarantees in AI systems (Richens & Everitt, 2024). While CRL can model complex
causal mechanisms, applying it to real-world environments with visual complexity remains challeng-
ing. Recent advancements in CRL (Lippe et al., 2022; 2023) have begun to tackle this problem in
simulated environments. These developments open up new possibilities for enhancing AI systems,
including LLMs. Although CRL does not directly address all LLM limitations, it can significantly
enhance their capabilities in specific domains. Our work builds upon these advancements, integrating
CRL with language models to improve their performance on causal inference and planning tasks.

We introduce a framework that combines CRL with language models to enable causally-aware
reasoning and planning in interactive environments. CRL provides LLMs with a structured, causal
understanding of the environment, reasoning about interventions and their consequences during
planning. The causal world model – akin to a simulator but learned rather than predefined – allows

1

000
001

002
003
004
005
006
007

008
009
010
011
012
013

014
015
016
017
018

019
020
021
022
023
024

025
026
027
028
029
030

031
032
033
034
035
036

037
038
039
040
041

042
043
044
045
046
047

048
049
050
051
052
053

Under review as a conference paper at ICLR 2025

054
055

056
057
058
059
060
061

062
063
064
065
066
067

068
069
070
071
072

073
074
075
076
077
078

079
080
081
082
083
084

085
086
087
088
089
090

091
092
093
094
095

096
097
098
099
100
101

102
103
104
105
106
107

Figure 1: Overview of a single rollout in the proposed planning pipeline. The causal encoder,
implemented using a CRL model, maps the high-dimensional state representation (image) to its
fundamental constituents—the causal variables. During planning, the LLM agent samples a proposed
action, which is then encoded by the text encoder. The causal transition model uses both the
disentangled latent representation of the image and the encoded action to simulate the next state
based on its learned causal mechanisms. The process iterates until the planning algorithm terminates,
with the causal model autoregressively operating in the latent space.

the LLM to evaluate multiple possible futures before taking action and thereby guides its decisions.
Conversely, LLMs offer a flexible interface for interacting with the causal world model, allowing for
more intuitive planning and reasoning that can leverage the LLM’s commonsense knowledge.

Furthermore, this work investigates using text to represent actions in the context of CRL-based world
modeling. Text-based action representations provide a flexible and intuitive way to describe actions,
making them more suitable for generalist agents operating in diverse environments. Moreover,
annotating frame sequences with natural language descriptions is often easier than exhaustively
enumerating every possible action in an environment, which can be intractable for complex domains.

We consider a setting with interleaved sequential observations in image format and corresponding
action descriptions at each timestep. This setup takes inspiration from real-world scenarios where
an agent might receive visual input (e.g., from a camera) along with descriptions of actions taken
(e.g., from system logs or human annotations). For example, in a robotic manipulation task, the
dataset might consist of a series of images showing the robot’s workspace, paired with descriptions
like “The gripper shifted slightly to the right.” or “The object was grasped and placed on the
worktop.” We assume no prior knowledge of the causal factors or the causal mechanisms between
them. The agent can only observe the effects of its actions from the images and does not require
explicit information about which specific variables or factors in the causal model it is affecting. Our
method builds on BISCUIT (Lippe et al., 2023), a CRL framework, to create a flexible causally-aware
world model from the sequence of observations and action descriptions, which is then used for
planning in environments.

The key contributions of our work are as follows:

• The first framework integrating CRL with LLMs to enable causally-aware reasoning and

planning in interactive environments.

• An exploration of text-based action representations for CRL and demonstration of their
effectiveness in data-scarce regimes, showing improved data efficiency in learning causal
representations.

• Demonstration of the framework’s effectiveness on a set of reasoning and planning tasks

across both static and dynamic environments.

Our experiments focus on simple environments, using existing CRL methods that are sufficiently
advanced for our use case. While these environments are still relatively simple, they represent the
current frontier of causal representation learning. As more powerful CRL methods become available,
they can be integrated into our framework, scaling it to more complex, realistic scenarios.

2 RELATED WORK

Causal Representation Learning Causal representation learning aims to identify the underlying
causal variables and their relations from high-dimensional observations (Schölkopf et al., 2021).
In the most general setting, the latent causal variables may not be uniquely identifiable (Locatello

2

CausalEncoderTextEncoderText DecoderToggle thethirdstoveknobPlace theplate downTextEncoderStopCausalTransitionModel......CausalTransitionModel.........LLMLLMLLMLLMText DecoderUnder review as a conference paper at ICLR 2025

108
109

110
111
112
113
114
115

116
117
118
119
120
121

122
123
124
125
126

127
128
129
130
131
132

133
134
135
136
137
138

139
140
141
142
143
144

145
146
147
148
149

150
151
152
153
154
155

156
157
158
159
160
161

et al., 2019a; Hyvärinen & Pajunen, 1999). Many approaches rely on assumptions or additional
knowledge about the causal structure, such as constraining the observation function (Buchholz et al.,
2023; Squires et al., 2023; Ahuja et al., 2023; Zhang et al., 2023; Kivva et al., 2022; Lachapelle et al.,
2023), sparse graphical structures (Khemakhem et al., 2020; Liu et al., 2022; 2024; Lachapelle &
Lacoste-Julien, 2022; Lachapelle et al., 2024), having multiple views (Xu et al., 2024; Yao et al.,
2024a; von Kügelgen et al., 2021; Brehmer et al., 2022; Locatello et al., 2020), or supplementary
supervision labels (Yang et al., 2020; Komanduri et al., 2022; Locatello et al., 2019b). Recent
advancements have explored CRL for temporal environments, in which agent-level actions like in
reinforcement learning are used to learn the causal structure of the environment (Lippe et al., 2022;
2023; Nalmpantis et al., 2023). In particular, our work leverages BISCUIT (Lippe et al., 2023), a
CRL framework that learns causal representations with realistic agent-focused assumptions, requiring
only a small set of labeled causal variables for the final mapping after causal representation learning,
without their interactions or causal graphs.

World Models and Causal Integration World models predict the consequences of actions and
have been extensively used in reinforcement learning (Ha & Schmidhuber, 2018). Recent work has
focused on object-centric world models (Greff et al., 2017; Steenkiste et al., 2018; Watters et al.,
2019) and the integration of graph neural networks for modeling transitions (Battaglia et al., 2016;
2018; Kipf et al., 2018). However, attempts to integrate causality into world models have been limited.
Some approaches, such as CoPhyNet (Baradel et al., 2020), consider counterfactual scenarios but
rely on direct supervision of object positions or place constraints on unobserved variables (Li et al.,
2020). Our work aims to learn a causal world model relying only on images and textual annotations 1
but capable of reasoning about actions across state transitions, while also being able to be interacted
with by a language model.

Large Language Models, Causality, Planning and Reasoning There has been much work
exploring the use of LLMs as action agents for planning and reasoning tasks, showing promising
results (Ahn et al., 2022; Hao et al., 2023; Huang et al., 2023). Various methodologies have been
developed to make use of LLMs for agent planning. These include task decomposition for breaking
complex tasks into subtasks (Wei et al., 2022; Yao et al., 2023; Shen et al., 2024), multi-plan selection
for generating and choosing optimal plans (Yao et al., 2024b; Wang et al., 2022), external module-
aided planning (Liu et al., 2023a; Guan et al., 2023), reflection and refinement via self-evaluation and
improvement (Shinn et al., 2024; Gou et al., 2024; Madaan et al., 2024), and memory-augmented
planning for decision making (Zhang et al., 2024; Zhong et al., 2024). While LLMs have shown
impressive performance in reasoning, tool usage, planning, and instruction-following, challenges
remain in addressing hallucinations, plan feasibility, and tractability in complex, multi-step planning
scenarios (Valmeekam et al., 2023; Kambhampati et al., 2024; Kambhampati, 2024). Theoretical
work on robustness under distribution shifts in unmediated decision tasks (where the decision does not
influence the utility) establishes a connection between causal understanding and robustness (Richens
& Everitt, 2024). A better approximation of the underlying causal model generally translates to more
robust agents, implying that world models should be causality-aware (Gupta et al., 2024).

3 BACKGROUND AND SETUP

To enable LLMs to perform causally-aware reasoning and planning in interactive environments, we
leverage CRL methods to build a causal world model (CWM). The CWM provides LLMs with a
structured understanding of the environment, allowing them to reason about interventions and their
consequences during planning.

In this section, we provide an overview of CRL in temporal causal graphs, which is foundational
to our framework. We discuss how CRL can learn latent causal representations from sequences of
observations and actions, setting the stage for integrating these representations with LLMs.

3.1 CAUSAL REPRESENTATION LEARNING IN TEMPORAL CAUSAL GRAPHS

CRL aims to uncover the latent causal variables and the underlying causal structure. In temporal
t=0, where Xt ∈ RD, and
settings, we consider sequences of high-dimensional observations {Xt}T

1Except for a few labels needed to map from the latent representation to human-interpretable language.

3

Under review as a conference paper at ICLR 2025

t=1, where Rt ∈ RE. Actions Rt can represent, for example, the coordinates of the
actions {Rt}T
locations where the interactions occurred (Lippe et al., 2023). The true causal variables {Ct}T
t=0,
where Ct ∈ RK, are unobserved. Furthermore, a deterministic observation model is assumed, often
represented as Xt = g(Ct), where g : RK → RD is an injective function mapping causal variables
to observations.

Instead of directly modeling causal variables, CRL relies on latent state representations. It estimates
a function f : RD → RM 2 that maps observations Xt to latent representations zt = f (Xt). The
goal is to ensure that each dimension zt
i in Ct up to a
transformation decided by the identifiability class of the causal model. Specifically, it aims to achieve
this disentanglement using only {Xt}T

i of zt corresponds to a causal variable C t

t=0 and {Rt}T

t=1.

3.2 GENERATIVE MODEL

The temporal CRL framework is often modeled as a generative process that describes how observa-
tions are produced from underlying latent state representations and actions. At each time step t, the
state zt evolve according to a transition model influenced by actions Rt, and generate observations
Xt. Assuming a first-order Markov process, the conditional likelihood of the observed data {Xt}T
t=0
given actions {Rt}T

t=1 is expressed as

p(cid:0){Xt} | {Rt}(cid:1) =

(cid:90)

p(z0)

T
(cid:89)

t=1

pω(zt | zt−1, Rt) pg(Xt | zt) dz,

(1)

where p(z0) is the prior distribution over the state. The transition model term pω(zt | zt−1, Rt)
models the state dynamics, capturing how the states evolve over time and how intervening actions
influence them. The observation model pg(Xt | zt) describes how the states generate the observations,
which in our case will be done with the deterministic function g.

The marginalization in Eq. (1) renders the objective intractable. A standard approach to address this is
to optimize the corresponding Evidence Lower Bound (ELBO) by assuming a Gaussian distribution
for the transition dynamics and the standard Gaussian for the prior, using the reparameterization trick
to enable efficient optimization (Kingma & Welling, 2013).

3.3

IDENTIFIABILITY GUARANTEES IN BISCUIT

There is nothing in the objective of Eq. (1) itself that guarantees the model will identify the causal
variables from the observations. In BISCUIT (Lippe et al., 2023), the CRL framework we adopt,
identifiability arises from two key assumptions: (1) each causal variable has a distinct ‘interaction
pattern,’ meaning that the effect of Rt on zt is mediated by a latent binary mask, and (2) these
interaction patterns vary over time. The first assumption is enforced by using a structured model
family to model the transition pω(zt | zt−1, Rt). We incorporate this component from BISCUIT in
our approach. These assumptions together ensure that causal variables are uniquely identifiable from
the observed data. For a more detailed discussion on the assumptions, theoretical guarantees, and the
structure of the transition model, we refer the reader to the original paper (Lippe et al., 2023).

4 BUILDING A CAUSAL WORLD MODEL FROM CAUSAL REPRESENTATIONS

To integrate the CRL model with LLMs, we construct a Causal World Model (CWM) that takes
actions in text format and states in image format and produces state representations in natural language.
The CWM builds on BISCUIT to model the environment’s dynamics, with the CWM’s encoder and
decoder components (see Figure 2) responsible for translating states and actions to and from natural
language. BISCUIT ensures identifiability and causal structure recovery, which enables reliable
predictions of the effects of actions/interventions, as demonstrated by our experiments in Section 6.
Although we utilize BISCUIT, our approach should be compatible with any CRL frameworks that can
provide disentangled causal representations (i.e., (Nalmpantis et al., 2023; Lachapelle et al., 2024;
Yao et al., 2022)).

2Since we do not normally have a priori information of the number of causal variables, we set M ≫ K.

4

162
163

164
165
166
167
168
169

170
171
172
173
174
175

176
177
178
179
180

181
182
183
184
185
186

187
188
189
190
191
192

193
194
195
196
197
198

199
200
201
202
203

204
205
206
207
208
209

210
211
212
213
214
215

Under review as a conference paper at ICLR 2025

Figure 2: Illustration of the first roll-out step with the Causal World Model. The image X0 and action
description L0 are encoded into initial latent representations. The CRL module then disentangles
these representations and the transition model predicts the next state. The causal mapper transforms
the disentangled causal representation of the next state into the estimated causal variables ˆC1. Finally,
the state descriptor s generates a natural language description ℓ1 of the next state. For subsequent
steps, the model can autoregress in the latent space using the previously predicted z, bypassing the
autoencoder and normalizing flow, enabling efficient multi-step inference and planning.

4.1 LANGUAGE GROUNDING MODULES

To integrate the CRL model with LLMs, we introduce architectural components that transform the
CRL model into a world model with a language interface. This section outlines the new components
we introduce, enabling the model to process image states and text inputs, and produce text outputs.

Language-Based Action Representations We replace the action encoding Rt in the CRL frame-
work with a language-based representation Le(Lt), where Le embeds a natural language description
Lt. This is implemented using an encoder-only language model (Reimers & Gurevych, 2019) with
a trainable head, replacing the original action encodings in the CRL framework’s transition model
Rt = Le(Lt) (see also Section 4.2).

Decoder The decoder G comprises two parts: the causal mapper and the state description generator.
The causal mapper mθ extracts causal variables C from the learned disentangled representations
z. It first identifies which latent dimensions zi are most predictive for each causal variable Cj, then
learns to perform the actual mapping. The state description generator s maps the estimated causal
variables ˆC to ℓ, a natural language description of the state. Detailed implementations of these
components are provided in Appendix G and H respectively.

4.2 PARAMETER ESTIMATION AND INFERENCE

In this section, we explain the estimation process for all model components and detail how the
resulting model is applied during inference. We use the GridWorld environment as a running example
to illustrate the process, though the same methodology applies to any environment.

Estimation: Causal Encoder and Transition Model To estimate the model, we use image pairs
{Xt} and corresponding action descriptions {Lt} in natural language, for example, “you toggled
the cyan traffic light” or “you moved the blue car”. We first train an autoencoder to compress high-
dimensional observations Xt into lower-dimensional latent representations Et = eψ(Xt), in which,
however, the causal variables will still be entangled. Then, analogously to Eq. 1, the conditional
likelihood of the encoded observations {Et}T

t=1 is given by

t=0 given action descriptions {Lt}T
T
(cid:89)

pω(zt | zt−1, Le(Lt)) pϕ(Et | zt) dz,

(2)

p(cid:0){Et} | {Lt}(cid:1) =

(cid:90)

p(z0)

t=1

where pω is the transition model, structured as in BISCUIT in order to satisfy the identifiability
guarantees, pϕ(Et | zt), is the observation model, and p(z0) is the prior distribution over the
initial latents, assumed to be the standard Gaussian. The invertible mapping fϕ : RM → RM is a
normalizing flow (NF) that transforms the autoencoder’s states {Et} into a new, structured latent
space {zt}, while identifying and separating the causal variables. As the NF is invertible, fϕ and its
Jacobian also yield the term pϕ(Et | zt) in the generative model. Similar to BISCUIT, the ELBO is
formulated and optimized using the reparameterization trick (Kingma & Welling, 2013).

5

216
217

218
219
220
221
222
223

224
225
226
227
228
229

230
231
232
233
234

235
236
237
238
239
240

241
242
243
244
245
246

247
248
249
250
251
252

253
254
255
256
257

258
259
260
261
262
263

264
265
266
267
268
269

You changed thestate of the cyantraffic light.CLIPTextEncoderAutoencoder...10.150.78...The cyan traffic lightis showing agreen signal. The redcar is positioned at(0.15, 0.75). ... Theheavy boulder ispositioned at (0.10,0.77).❄❄StateDescriptorMLPMLPMLP...Causal MapperTargetAssignmentNF0101CRL ModelCausal EncoderInputDecoderCausal TransitionModel......MLPBISCUITUnder review as a conference paper at ICLR 2025

270
271

272
273
274
275
276
277

278
279
280
281
282
283

284
285
286
287
288

289
290
291
292
293
294

295
296
297
298
299
300

301
302
303
304
305
306

307
308
309
310
311

312
313
314
315
316
317

318
319
320
321
322
323

The full causal encoder that maps an observation Xt to the causal state zt is expressed as E := fϕ ◦eψ,
where eψ is the encoder part of the autoencoder, and fϕ is the NF. While we assume perfect
reconstruction capability for the autoencoder, a common assumption in CRL (Kivva et al., 2022;
Lachapelle et al., 2023; Brehmer et al., 2022; Lachapelle et al., 2024, inter alia), this component can
be replaced with stronger visual encoders as they become available, without affecting the framework’s
core functionality.

This framework builds upon the BISCUIT architecture, maintaining the same structure for the
autoencoder, normalizing flow (RealNVP (Dinh et al., 2016)), and transition model. However, we
introduce an important modification in the action representation. While BISCUIT used coordinate-
based action encodings (as described in Section 5.1), our work incorporates language-based action
representations through Le(Lt) in the transition model to enable our model to process natural
language action descriptions.

Estimation: Decoder We train the causal mapper mθ using a small set of annotated but not
necessarily ordered images where the true causal variables C and their values are known. The
training pairs consist of (z, C), where z is the output of our causal encoding pipeline and C are
the corresponding ground truth causal variables. In GridWorld, these variables include positions of
vehicles and obstacles, and states of traffic lights. For instance, the causal mapper might learn that
dimensions z1, z3, and z7 are most predictive for the “blue car x-position” (C0), and then train a
specific predictor for C0 using only these relevant dimensions.

The state description generator s, typically a rule-based system, maps the estimated causal variables
to human-interpretable outputs. For example, it might transform position and state variables into a
description like “The blue car is at (2,3), the cyan traffic light is green”. The full decoder is expressed
as G := s ◦ mθ.

Inference Process During inference, the model sequentially processes new GridWorld images
through these components:

zt = E(Xt) = (fϕ ◦ eψ)(Xt),
(cid:0)zt+1 | zt, Le(Lt)(cid:1),
zt+1 ∼ pω
ℓt+1 = G(zt+1) = (s ◦ mθ)(zt+1).
This process transforms raw input into interpretable state descriptions of the next state, facilitating
interaction with language models for reasoning and planning tasks. Notably, the transition model
operates solely in the disentangled latent space, without dependency on the high-dimensional observa-
tions Xt. This enables efficient multi-step inference through autoregression, allowing for long-term
planning and reasoning without the need to decode back to the observation space at each step.

This entire process relies solely on the sequence of observations and action descriptions, without
requiring explicit information about which specific variables or factors are being affected. By
introducing the language-based action encoder and decoding into natural language, we create a
framework that is inherently suited for language-based causal reasoning in complex environments.
The algorithm to perform inference is provided in Appendix M.

5 EXPERIMENTAL SETUP

We evaluate our framework using two distinct environments: a dynamic 8 × 8 GridWorld and a static
3D-rendered kitchen (AI2-THOR) (Kolve et al., 2017). The GridWorld is dynamic, meaning the
environment state can change even without agent actions, while the iTHOR kitchen is static, changing
only in response to agent interventions. Our experiments focus on three key aspects: the effectiveness
of text-based action representations, causal inference, and planning. Both environments feature
various objects with causal variables representing their states and positions. Detailed descriptions of
the environments are provided in Appendices A and B.

For each environment, we generated multiple datasets for training, evaluation, and in-context learning.
The data generation process involves initializing the environment state and performing random,
valid actions. Specific details about dataset sizes, in-context learning example generation, and
self-evaluation reward generation for planning tasks are described in Appendix D.

6

Under review as a conference paper at ICLR 2025

324
325

326
327
328
329
330
331

332
333
334
335
336
337

338
339
340
341
342

343
344
345
346
347
348

349
350
351
352
353
354

355
356
357
358
359
360

361
362
363
364
365

366
367
368
369
370
371

372
373
374
375
376
377

5.1 ACTION REPRESENTATIONS

We investigate three action representation modalities:

1. Coordinate-based (CB): Encoding of 2D pixel coordinates indicating the position where
the interaction was performed. For example, a click at position (2, 3) is transformed into a
higher-dimensional representation using high-frequency sinusoidal functions.

2. Text-based (TB): Natural language descriptions expanded using a PCFG (e.g., “you toggled

the bright cyan traffic light”), then encoded through an encoder-only text embedder.

3. Hybrid (HB): Concatenation of coordinate-based and text-based representations.

We hypothesize that the text-based action encoding is a) semantically richer, providing more infor-
mation for the same or less effort to annotate the data, b) more flexible, enabling a language-based
interface suitable for a generalist agent, and c) more robust, meaning that paraphrases or equivalent
descriptions of the same action can still work with our model even if it was not specifically trained on
them. This last point is crucial, as the LLM used at inference may deviate in its action description
style from what was seen during training.

5.2 BASELINE

Our baseline uses the world model component from the Reasoning via Planning (RAP) methodology
(Hao et al., 2023). This language model-based world model predicts the next state given the current
state st, chosen action at, and context c:

st+1 ∼ pLM(st+1 | st, at, c).
The baseline constructs a prompt at runtime that includes the environment description and dynamics,
current state representation, chosen action, two relevant in-context learning (ICL) examples, and
instructions for predicting the next state. This approach leverages the language model’s pretrained
knowledge while adapting to the specific task and environment dynamics. We ensure the relevance of
the ICL examples by providing examples that match the current action and the object it is applied to.

We use LLaMA 3 (8B) (Dubey et al., 2024) as the planning agent quantized to 6 bits in the exl2
format. We chose RAP+LLaMA3 as the baseline for its simplicity and effectiveness, providing a fair
point of comparison to assess the benefits of integrating causal representation learning. This allows
us to isolate the impact of causal understanding in an otherwise comparable framework, though our
approach could integrate with alternative search algorithms such as LLM-MCTS (Zhao et al., 2024).

6 EXPERIMENTS AND DISCUSSION

6.1 EVALUATION OF TEXT-BASED ACTION REPRESENTATIONS

In this experiment, we demonstrate the effectiveness of representing actions in natural language
for learning causal representations. We assess the induced state variables z by comparing them to
ground-truth causal variables. Note that the model’s decoder is not evaluated in these experiments.

We train our causal world model using each action modality (CB, TB, HB) across different sub-
sample percentages of the training dataset, focusing on the low-data regime. Given sufficient data,
models yield practically identical results across all 3 modalities but obtaining data in non-simulated
environments is typically expensive. Performance is assessed using a standard CRL metric: R2
scores for the permutation π that maximizes the diagonal of the R2 matrix between learned latent
variables and true causal variables. This approach accounts for the fact that we learn causal variables
up to permutation. Each experiment uses 3 seeds with distinct subsamples. A more comprehensive
explanation of the training of the components of the CRL models used is presented in Appendix E.

Table 1 presents the results of our action representation experiments for the GridWorld environment in
the low-data regime. Our results demonstrate that incorporating text into action representations (TB
and HB) is Pareto-optimal in GridWorld; TB and HB perform at least as well as the coordinate-based
representation, especially in low-data regimes. In extremely low-data scenarios (0.3% - 0.7%), the
hybrid approach consistently outperforms both CB and TB. As data increases (1.0% - 1.5%), TB
shows competitive performance while providing natural alignment with LLM interfaces.

7

Under review as a conference paper at ICLR 2025

378
379

380
381
382
383
384
385

386
387
388
389
390
391

392
393
394
395
396

397
398
399
400
401
402

403
404
405
406
407
408

409
410
411
412
413
414

415
416
417
418
419

420
421
422
423
424
425

426
427
428
429
430
431

Table 1: R2 scores for action representations. CB: Coordinate-based, TB: Text-based, HB: Hybrid.
100% stands for 106 image states.

Action
Type

0.3%

0.5%

0.7%

1.0%

1.2%

1.5%

Subsample Percentage

100%

CB
TB
HB

0.392 ±0.000 0.366±0.000 0.424±0.001 0.457±0.004 0.472±0.004 0.548±0.022 0.987
0.374±0.000 0.362±0.000 0.399±0.000 0.470±0.012 0.495±0.014 0.603±0.003 0.990
0.392±0.000 0.433±0.001 0.460±0.000 0.461±0.007 0.490±0.010 0.539±0.011 0.991

Table 2: N -step causal inference accuracies for the causal world model and the RAP (Hao et al.,
2023) world model across different environments and step lengths.

Steps

iTHOR
2

1

4

1

GridWorld
2

4

6

8

0.824 0.680 0.630 0.954 0.922 0.829 0.797 0.758
Causal Model
RAP World Model 0.482 0.285 0.110 0.391 0.220 0.085 0.045 0.005

These findings support our hypothesis: action encodings including text are as effective as or superior
in uncovering causal variables to coordinate-based representations, particularly when data is scarce.
We use the TB model trained using the entire dataset (106 examples) in the subsequent experiments.

6.2 CAUSAL INFERENCE PERFORMANCE

Our causal inference experiments evaluate both world models’ ability to perform 1-step and N -step
causal inference, i.e., predict the effects of actions (interventions) on the environment. In the 1-step
case, given the current state and an action, the model predicts the new state. For N -step causal
inference, we provide a sequence of actions and only the starting state and the world model applies
each action to its previous prediction in a sequence. This differs from planning in that it focuses on
the effect of a given sequence of actions rather than finding actions to reach a goal. The evaluation
methodology is presented in Appendix K.

Table 2 presents accuracies of causal inference for both models across different environments and
step lengths. The causal world model consistently outperforms the baseline across all scenarios.
In GridWorld, it maintains high accuracy (75.8%) even for 8-step inference, while the baseline’s
performance drops nearly to 0. The performance in iTHOR, while lower than in GridWorld, still
shows a substantial improvement over the baseline.

The higher overall performance on GridWorld can be attributed to its simpler action space, object
space, and causal graph, despite its dynamic nature. The baseline’s lower performance in GridWorld
compared to iTHOR may be due to the lack of visual input, which is less natural for language models
in an artificial environment.

Table 3 provides a detailed breakdown of the causal inference performance for specific actions and
objects, based on the extended 1-step dataset of 3000 samples. In iTHOR, the causal world model
excels at ToggleObject and OpenObject actions (95.7% and 92.6% accuracy), while struggling more
with PutObject and PickupObject actions (50.6% and 43.1% accuracy). This discrepancy likely
stems from the following; first, we model the three-dimensional coordinates as independent random
variables while, in reality, they are dependent. Second, we model interventions using binary variables
to estimate whether we performed an intervention or not. Performance could be improved by injecting
inductive bias towards the continuous, three-dimensional nature of the underlying variable. However,
this requires task specialization within the model and we chose to keep the proposed framework task-
agnostic. The baseline model shows a different pattern, performing better on NoOp and PickupObject
actions but struggling with PutObject actions.

In GridWorld, the causal world model demonstrates high accuracy across all action types, with
particularly strong performance in changes to the state of the lights and no-action scenarios. The

8

Under review as a conference paper at ICLR 2025

432
433

434
435
436
437
438
439

440
441
442
443
444
445

446
447
448
449
450

451
452
453
454
455
456

457
458
459
460
461
462

463
464
465
466
467
468

469
470
471
472
473

474
475
476
477
478
479

480
481
482
483
484
485

Table 3: Causal inference accuracies for action categories in iTHOR and GridWorld environments.
CWM: Causal World Model, RAP (Hao et al., 2023): RAP World Model Baseline.

iTHOR Environment

GridWorld Environment

Action Category CWM RAP Action Category

CWM RAP

ToggleObject
OpenObject
NoOp
PutObject
PickupObject

0.957
0.926
0.962
0.506
0.431

0.466 Change Light State
0.339 No Action
0.710 Move
0.100
0.692

0.986
0.985
0.928

0.300
0.456
0.408

baseline model shows lower performance across the board, with its best performance on the No
Action category.

These results highlight the causal world model’s superior ability to reason about causal relationships,
maintaining strong performance across different temporal scales, environments, and action types.

6.3 PLANNING

Methodology The planning experiments assess the model’s ability to generate a sequence of
actions to transform an initial state into a goal state. This involves exploring multiple possible action
sequences and evaluating their effectiveness in reaching the goal. Unlike causal inference, planning
requires considering long-term consequences and optimizing for a specific objective.

Our framework adapts the Reasoning via Planning (RAP) methodology, with a key distinction: we
employ a separate causal world model alongside a language model agent, rather than using a single
language model for both roles. We use the same LLM as for the baseline planning agent (LLaMA 3).

The planning works as follows: The LLaMA 3 agent proposes possible actions based on the current
state. The world model then simulates the actions’ outcomes, predicting subsequent states. The agent
then evaluates each state-action pair’s quality and picks an action resulting in a new state. This cycle
repeats, exploring multiple reasoning paths before converging on a final solution. For all N -step
experiments, we use a search tree depth of N + 2. We use a modified version of the RAP-MCTS
algorithm, presented in Appendix N.

In Gridworld, there are three actions to toggle traffic lights (one per light) and one to
Actions
perform no action. In iTHOR, we dynamically generate 10-15 possible actions, depending on the
initial state. During planning, the models use their internal representations to determine possible
actions. During evaluation, we use the external simulator (the same one used to generate the data) to
execute the plan proposed by the agent. If an invalid action is proposed, during evaluation we default
to performing no action.

Reward Design In line with the RAP methodology, we rely on the LLM’s ability to judge the
current state in relation to the goal. The Intuition reward is the unnormalized log probability of
actions generated by the language model, given the current state and few-shot demonstrations. The
Self-evaluation reward is the log probability of the token “good” when asking the model to evaluate
whether the proposed action is correct, given the current state and few-shot demonstrations.

We avoid using percentage-of-goals-reached rewards to maintain generality and applicability to
problems that are not easily divisible into subgoals or subtasks. This choice ensures that our method
remains applicable to a wide range of problems, including those where intermediate progress toward
the goal is difficult to quantify and/or may not correlate directly with overall success.

6.3.1 PLANNING RESULTS AND DISCUSSION

Table 4 presents the planning results for both models across different environments and step lengths.

The causal world model consistently outperforms the baseline in both environments:

9

Under review as a conference paper at ICLR 2025

486
487

488
489
490
491
492
493

494
495
496
497
498
499

500
501
502
503
504

505
506
507
508
509
510

511
512
513
514
515
516

517
518
519
520
521
522

523
524
525
526
527

528
529
530
531
532
533

534
535
536
537
538
539

Table 4: Planning results for the causal world model and RAP baseline (Hao et al., 2023) across
different environments and step lengths. The best performing method for each metric is highlighted
in green.

Environment

Steps

Causal World Model

RAP Baseline

Success Rate ↑

Avg. Steps Avg. Steps
(Failure) ↓
(Success) ↓

Success Rate ↑

Avg. Steps Avg. Steps
(Failure) ↓
(Success) ↓

iTHOR

GridWorld

2
4

2
4
6
8

0.58
0.44

0.95
0.73
0.46
0.42

1.78
2.14

1.92
2.71
3.65
4.62

3.36
5.43

3.20
5.27
7.72
9.76

0.25
0.11

0.20
0.11
0.08
0.06

4.00
6.00

2.00
3.27
5.50
7.00

4.00
6.00

3.19
5.48
7.72
9.93

• Success Rates: The causal model achieves significantly higher success rates, particularly in
longer planning horizons. In iTHOR, it more than doubles the baseline’s success rate for
2-step planning (0.58 vs 0.25) and quadruples it for 4-step planning (0.44 vs 0.11).

• Efficiency: For successful trajectories, the causal model takes fewer steps on average to

reach the goal state, indicating more efficient planning.

• Scalability: While both models show decreased performance as the number of steps in
the ground truth increase, the causal model degrades more gracefully. In GridWorld, it
maintains a 0.42 success rate for 8-step planning, compared to the baseline’s 0.06.

• Consistency: Both models perform better in GridWorld compared to iTHOR, likely due to
the lower complexity and more constrained action space. However, the causal model shows
more consistent performance across both environments.

An interesting observation is the sub-N performance in N -step planning scenarios. This phenomenon
arises from two key factors in our experimental design which renders the parameter N an upper
bound of the steps needed to achieve the goal state. In the static iTHOR environment, some actions
can negate others (e.g., toggling the toaster twice is equivalent to performing no action), allowing for
shorter paths to the goal state. In addition to this phenomenon, in the dynamic GridWorld environment,
the inherent movement of entities (e.g., cars moving when facing a green light) can sometimes lead
to the goal state in fewer steps than the upper bound. This sub-N performance highlights our models’
ability to find efficient paths to the goal state, often outperforming the original trajectories used to
generate the planning problems.

The performance improvements observed in our experiments can be attributed to the CRL world
model’s higher accuracy in predicting future states. Both our method and the baseline use identical
state representations, with consistent text formatting maintained across the in-context learning
examples. This uniformity ensures that the language model’s reasoning and planning processes are
influenced primarily by the accuracy of the underlying world model. Therefore, the superior planning
performance of our method highlights the effectiveness of integrating CRL for more accurate state
predictions, which directly benefits the downstream reasoning tasks.

7 CONCLUSION

In this work, we introduced a framework that integrates causal representation learning with language
models, enabling causally-aware reasoning and planning in interactive environments. Our approach
combines the structured causal understanding of CRL with the flexible interface of language models,
demonstrating superior performance in causal inference and planning tasks across two environments.
The causal world model consistently outperforms baselines, showing improved accuracy, efficiency,
and scalability as task complexity increases. Our exploration of text-based action representations
reveals potential advantages in low-data regimes, suggesting implications for more flexible and
generalizable AI systems. While our current experiments focus on relatively simple environments,
the framework is designed to extend to more complex scenarios as CRL and search methods advance.
Future work could explore applications to real-world environments, improve the interpretability of
learned causal world models, and develop techniques independent of labeled causal variables.

10

Under review as a conference paper at ICLR 2025

540
541

542
543
544
545
546
547

548
549
550
551
552
553

554
555
556
557
558

559
560
561
562
563
564

565
566
567
568
569
570

571
572
573
574
575
576

577
578
579
580
581

582
583
584
585
586
587

588
589
590
591
592
593

REFERENCES

Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea
Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Daniel Ho, Jasmine
Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Eric Jang, Rosario Jauregui Ruano, Kyle Jeffrey,
Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Kuang-Huei Lee,
Sergey Levine, Yao Lu, Linda Luu, Carolina Parada, Peter Pastor, Jornell Quiambao, Kanishka
Rao, Jarek Rettinghouse, Diego Reyes, Pierre Sermanet, Nicolas Sievers, Clayton Tan, Alexander
Toshev, Vincent Vanhoucke, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Mengyuan Yan, and Andy
Zeng. Do as i can and not as i say: Grounding language in robotic affordances. In arXiv preprint
arXiv:2204.01691, 2022.

Kartik Ahuja, Divyat Mahajan, Yixin Wang, and Yoshua Bengio. Interventional Causal Representation
Learning. In Proceedings of the 40th International Conference on Machine Learning (ICML),
Proceedings of Machine Learning Research. PMLR, PMLR, 2023.

Fabien Baradel, Natalia Neverova, Julien Mille, Greg Mori, and Christian Wolf. Cophy: Counterfac-

tual learning of physical dynamics. In ICLR, 2020.

Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks
for learning about objects, relations and physics. Advances in neural information processing
systems, 29, 2016.

Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi,
Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al.
Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261,
2018.

Johann Brehmer, Pim de Haan, Phillip Lippe, and Taco Cohen. Weakly Supervised Causal Represen-
tation Learning. In Advances in Neural Information Processing Systems (NeurIPS), volume 35.
Curran Associates, Inc., 2022.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are
few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.

Simon Buchholz, Goutham Rajendran, Elan Rosenfeld, Bryon Aragam, Bernhard Schölkopf, and
Pradeep Kumar Ravikumar. Learning Linear Causal Representations From Interventions Under
General Nonlinear Mixing. In Advances in Neural Information Processing Systems (NeurIPS),
volume 36. Curran Associates, Inc., 2023.

Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International

conference on computers and games, pp. 72–83. Springer, 2006.

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv

preprint arXiv:1605.08803, 2016.

Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha
Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models.
arXiv preprint arXiv:2407.21783, 2024.

Zhibin Gou, Zhihong Shao, Yeyun Gong, Yelong Shen, Yujiu Yang, Nan Duan, and Weizhu Chen.
CRITIC: large language models can self-correct with tool-interactive critiquing. In The Twelfth In-
ternational Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024.
OpenReview.net, 2024. URL https://openreview.net/forum?id=Sx038qxjek.

Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization.

Advances in Neural Information Processing Systems, 30, 2017.

Lin Guan, Karthik Valmeekam, Sarath Sreedharan, and Subbarao Kambhampati. Leveraging pre-
trained large language models to construct and utilize world models for model-based task planning.
Advances in Neural Information Processing Systems, 36:79081–79094, 2023.

11

Under review as a conference paper at ICLR 2025

594
595

596
597
598
599
600
601

602
603
604
605
606
607

608
609
610
611
612

613
614
615
616
617
618

619
620
621
622
623
624

625
626
627
628
629
630

631
632
633
634
635

636
637
638
639
640
641

642
643
644
645
646
647

Tarun Gupta, Wenbo Gong, Chao Ma, Nick Pawlowski, Agrin Hilmkil, Meyer Scetbon, Ade Famoti,
Ashley Juan Llorens, Jianfeng Gao, Stefan Bauer, et al. The essential role of causality in foundation
world models for embodied ai. arXiv preprint arXiv:2402.06665, 2024.

David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122, 2018.

Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu.
Reasoning with language model is planning with world model. arXiv preprint arXiv:2305.14992,
2023.

Wenlong Huang, Chen Wang, Ruohan Zhang, Yunzhu Li, Jiajun Wu, and Li Fei-Fei. Voxposer:
Composable 3d value maps for robotic manipulation with language models. arXiv preprint
arXiv:2307.05973, 2023.

Aapo Hyvärinen and Petteri Pajunen. Nonlinear Independent Component Analysis: Existence and

Uniqueness Results. Neural Networks, 12(3):429–439, 1999.

Subbarao Kambhampati. Can large language models reason and plan? Annals of the New York

Academy of Sciences, 1534(1):15–18, 2024.

Subbarao Kambhampati, Karthik Valmeekam, Lin Guan, Mudit Verma, Kaya Stechly, Siddhant
Bhambri, Lucas Paul Saldyt, and Anil B Murthy. Position: LLMs can’t plan, but can help planning
in LLM-modulo frameworks. In Forty-first International Conference on Machine Learning, 2024.
URL https://openreview.net/forum?id=Th8JPEmH4z.

Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational Autoencoders
and Nonlinear ICA: A Unifying Framework. In Proceedings of the 23rd International Conference
on Artificial Intelligence and Statistics (AISTATS), Proceedings of Machine Learning Research.
PMLR, 2020.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes.

arXiv preprint

arXiv:1312.6114, 2013.

Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational
inference for interacting systems. In International conference on machine learning, pp. 2688–2697.
PMLR, 2018.

Bohdan Kivva, Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Identifiability of Deep
Generative Models Without Auxiliary Information. In Advances in Neural Information Processing
Systems (NeurIPS), volume 35. Curran Associates, Inc., 2022.

Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In Johannes Fürnkranz,
Tobias Scheffer, and Myra Spiliopoulou (eds.), Machine Learning: ECML 2006, pp. 282–293,
Berlin, Heidelberg, 2006. Springer Berlin Heidelberg. ISBN 978-3-540-46056-5.

Eric Kolve, Roozbeh Mottaghi, Winson Han, Eli VanderBilt, Luca Weihs, Alvaro Herrasti, Daniel
Gordon, Yuke Zhu, Abhinav Gupta, and Ali Farhadi. AI2-THOR: An Interactive 3D Environment
for Visual AI. arXiv, 2017.

Aneesh Komanduri, Yongkai Wu, Wen Huang, Feng Chen, and Xintao Wu. Scm-vae: Learning
identifiable causal representations via structural knowledge. In 2022 IEEE International Conference
on Big Data (Big Data), pp. 1014–1023. IEEE, 2022.

Sébastien Lachapelle and Simon Lacoste-Julien. Partial Disentanglement via Mechanism Sparsity.

In UAI 2022 Workshop on Causal Representation Learning, 2022.

Sebastien Lachapelle, Divyat Mahajan, Ioannis Mitliagkas, and Simon Lacoste-Julien. Additive
Decoders for Latent Variables Identification and Cartesian-Product Extrapolation. In Advances in
Neural Information Processing Systems (NeurIPS), volume 36. Curran Associates, Inc., 2023.

Sébastien Lachapelle, Pau Rodríguez López, Yash Sharma, Katie Everett, Rémi Le Priol, Alexandre
Lacoste, and Simon Lacoste-Julien. Nonparametric Partial Disentanglement via Mechanism
Sparsity: Sparse Actions, Interventions and Sparse Temporal Dependencies. arXiv preprint
arXiv:2401.04890, 2024.

12

Under review as a conference paper at ICLR 2025

648
649

650
651
652
653
654
655

656
657
658
659
660
661

662
663
664
665
666

667
668
669
670
671
672

673
674
675
676
677
678

679
680
681
682
683
684

685
686
687
688
689

690
691
692
693
694
695

696
697
698
699
700
701

Minne Li, Mengyue Yang, Furui Liu, Xu Chen, Zhitang Chen, and Jun Wang. Causal world models
by unsupervised deconfounding of physical dynamics. arXiv preprint arXiv:2012.14228, 2020.

Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Stratis Gavves. Citris:
Causal identifiability from temporal intervened sequences. In International Conference on Machine
Learning, pp. 13557–13603. PMLR, 2022.

Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves.
Biscuit: Causal representation learning from binary interactions. In Uncertainty in Artificial
Intelligence, pp. 1263–1273. PMLR, 2023.

Bo Liu, Yuqian Jiang, Xiaohan Zhang, Qiang Liu, Shiqi Zhang, Joydeep Biswas, and Peter Stone.
Llm+p: Empowering large language models with optimal planning proficiency. arXiv preprint
arXiv:2304.11477, 2023a.

Rosanne Liu, Joel Lehman, Piero Molino, Felipe Petroski Such, Eric Frank, Alex Sergeev, and
Jason Yosinski. An intriguing failing of convolutional neural networks and the coordconv solution.
Advances in neural information processing systems, 31, 2018.

Yuhang Liu, Zhen Zhang, Dong Gong, Mingming Gong, Biwei Huang, Anton van den Hengel, Kun
Zhang, and Javen Qinfeng Shi. Identifying Weight-Variant Latent Causal Models. arXiv preprint
arXiv:2208.14153, 2022.

Yuhang Liu, Zhen Zhang, Dong Gong, Mingming Gong, Biwei Huang, Anton van den Hengel, Kun
Zhang, and Javen Qinfeng Shi. Identifiable Latent Polynomial Causal Models Through the Lens of
Change. In Proceedings of the 12th International Conference on Learning Representations (ICLR),
2024.

Zhihan Liu, Hao Hu, Shenao Zhang, Hongyi Guo, Shuqi Ke, Boyi Liu, and Zhaoran Wang. Reason
for future, act for now: A principled framework for autonomous llm agents with provable sample
efficiency. arXiv preprint arXiv:2309.17382, 2023b.

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf,
and Olivier Bachem. Challenging Common Assumptions in the Unsupervised Learning of Dis-
entangled Representations. In Proceedings of the 36th International Conference on Machine
Learning (ICML), Proceedings of Machine Learning Research. PMLR, 2019a.

Francesco Locatello, Michael Tschannen, Stefan Bauer, Gunnar Rätsch, Bernhard Schölkopf,
and Olivier Bachem. Disentangling factors of variation using few labels. arXiv preprint
arXiv:1905.01258, 2019b.

Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael
Tschannen. Weakly-Supervised Disentanglement Without Compromises. In Proceedings of the
37th International Conference on Machine Learning (ICML), Proceedings of Machine Learning
Research. PMLR, 2020.

Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri
Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, et al. Self-refine: Iterative refinement
with self-feedback. Advances in Neural Information Processing Systems, 36, 2024.

Angelos Nalmpantis, Phillip Lippe, and Sara Magliacane. Hierarchical Causal Representation

Learning. In Causal Representation Learning Workshop at NeurIPS 2023, 2023.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language

models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.

Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks.

arXiv preprint arXiv:1908.10084, 2019.

Jonathan Richens and Tom Everitt. Robust agents learn causal world models.

In The Twelfth
International Conference on Learning Representations, 2024. URL https://openreview.
net/forum?id=pOoKI3ouv1.

13

Under review as a conference paper at ICLR 2025

702
703

704
705
706
707
708
709

710
711
712
713
714
715

716
717
718
719
720

721
722
723
724
725
726

727
728
729
730
731
732

733
734
735
736
737
738

739
740
741
742
743

744
745
746
747
748
749

750
751
752
753
754
755

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, 109(5):612–634, 2021.

Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu, and Yueting Zhuang. Hugginggpt:
Solving ai tasks with chatgpt and its friends in hugging face. Advances in Neural Information
Processing Systems, 36, 2024.

Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion:
Language agents with verbal reinforcement learning. Advances in Neural Information Processing
Systems, 36, 2024.

Chandler Squires, Anna Seigal, Salil S. Bhate, and Caroline Uhler. Linear Causal Disentanglement
via Interventions. In Proceedings of the 40th International Conference on Machine Learning
(ICML), Proceedings of Machine Learning Research. PMLR, PMLR, 2023.

Sjoerd Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural expecta-

tion maximization: Unsupervised discovery of objects and their interactions. 02 2018.

Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the
planning abilities of large language models-a critical investigation. Advances in Neural Information
Processing Systems, 36:75993–76005, 2023.

Julius von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel
Besserve, and Francesco Locatello. Self-Supervised Learning With Data Augmentations Provably
Isolates Content From Style. In Advances in Neural Information Processing Systems (NeurIPS),
volume 34. Curran Associates, Inc., 2021.

Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdh-
ery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models.
arXiv preprint arXiv:2203.11171, 2022.

Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner.
Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven
exploration. arXiv preprint arXiv:1905.09275, 2019.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny
Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in
neural information processing systems, 35:24824–24837, 2022.

Danru Xu, Dingling Yao, Sebastien Lachapelle, Perouz Taslakian, Julius von Kügelgen, Francesco
Locatello, and Sara Magliacane. A Sparsity Principle for Partially Observable Causal Representa-
tion Learning. In Proceedings of the 41st International Conference on Machine Learning (ICML),
Proceedings of Machine Learning Research. PMLR, 2024.

Mengyue Yang, Furui Liu, Zhitang Chen, Xinwei Shen, Jianye Hao, and Jun Wang. Causalvae:
Structured causal disentanglement in variational autoencoder. arXiv preprint arXiv:2004.08697,
2020.

Dingling Yao, Danru Xu, Sebastien Lachapelle, Sara Magliacane, Perouz Taslakian, Georg Martius,
Julius von Kügelgen, and Francesco Locatello. Multi-View Causal Representation Learning
With Partial Observability. In Proceedings of the 12th International Conference on Learning
Representations (ICLR), 2024a.

Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao.
ReAct: Synergizing reasoning and acting in language models. In International Conference on
Learning Representations (ICLR), 2023.

Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan.
Tree of thoughts: Deliberate problem solving with large language models. Advances in Neural
Information Processing Systems, 36, 2024b.

14

Under review as a conference paper at ICLR 2025

Weiran Yao, Guangyi Chen, and Kun Zhang. Temporally Disentangled Representation Learning. In
Advances in Neural Information Processing Systems (NeurIPS), volume 35. Curran Associates,
Inc., 2022.

Matej Zeˇcevi´c, Moritz Willig, Devendra Singh Dhami, and Kristian Kersting. Causal parrots: Large
language models may talk causality but are not causal. arXiv preprint arXiv:2308.13067, 2023.

Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language
image pre-training. In Proceedings of the IEEE/CVF International Conference on Computer Vision,
pp. 11975–11986, 2023.

Danyang Zhang, Lu Chen, Situo Zhang, Hongshen Xu, Zihan Zhao, and Kai Yu. Large language
models are semi-parametric reinforcement learning agents. Advances in Neural Information
Processing Systems, 36, 2024.

Jiaqi Zhang, Kristjan Greenewald, Chandler Squires, Akash Srivastava, Karthikeyan Shanmugam,
and Caroline Uhler. Identifiability Guarantees for Causal Disentanglement From Soft Interventions.
In Advances in Neural Information Processing Systems (NeurIPS), volume 36. Curran Associates,
Inc., 2023.

Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models as commonsense knowledge for

large-scale task planning. Advances in Neural Information Processing Systems, 36, 2024.

Wanjun Zhong, Lianghong Guo, Qiqi Gao, He Ye, and Yanlin Wang. Memorybank: Enhancing large
language models with long-term memory. In Proceedings of the AAAI Conference on Artificial
Intelligence, volume 38, pp. 19724–19731, 2024.

15

756
757

758
759
760
761
762
763

764
765
766
767
768
769

770
771
772
773
774

775
776
777
778
779
780

781
782
783
784
785
786

787
788
789
790
791
792

793
794
795
796
797

798
799
800
801
802
803

804
805
806
807
808
809

Under review as a conference paper at ICLR 2025

A GRIDWORLD ENVIRONMENT

The gridworld environment is a dynamic environment of size H × H, where H ∈ N denotes both the
height and width of the grid. The top left corner of the grid is defined to be (0, 0). The environment
consists of C underlying causal variables that interact based on the actions taken by the agent
and the dynamics of the environment. The environment contains three types of entities: vehicles
v ∈ V , obstacles o ∈ O, and traffic lights tl ∈ T L. Each entity has a fixed corresponding attribute,
implemented as a color, which differentiates it from other objects within the same entity class.

The traffic lights are positioned in the grid, and each vehicle is facing a specific traffic light. The
positions of the traffic lights are fixed and immutable, with coordinates (xtl, ytl), where xtl, ytl ∈
{0, 1, . . . , H − 1}. Each traffic light has a state stl ∈ {red, green}. The obstacles have positions
(xo, yo) in the grid, where xo, yo ∈ {0, 1, . . . , H − 1}, and these positions can only change through
interventions performed on them. The vehicles have positions (xv, yv) in the grid, where xv, yv ∈
{0, 1, . . . , H − 1}, and an orientation θv ∈ {up, down, left, right}. The vehicle positions change
according to the following dynamics:

Let v be a vehicle at position (xv, yv) with orientation θv, associated with a traffic light tl at position
(xtl, ytl). We say that the vehicle v is facing the traffic light tl if and only if one of the following
conditions is satisfied:

1. θv = up and xv = xtl and yv > ytl
2. θv = down and xv = xtl and yv < ytl
3. θv = left and yv = ytl and xv > xtl
4. θv = right and yv = ytl and xv < xtl

If the vehicle v is facing the traffic light tl, it will move forward to the cell (x′
timestep if and only if all of the following conditions are satisfied:

v, y′

v) at the next

1. The traffic light tl has a state of green, i.e., stl = green.
v, y′
2. There are no obstacles in the cell (x′
3. There are no traffic lights in the cell (x′
4. The cell (x′

v) is within the grid boundaries, i.e., 0 ≤ x′

v, y′

v), i.e., ∄ o ∈ O : (xo, yo) = (x′
v, y′

v).
v), i.e., ∄ tl ∈ T L : (xtl, ytl) = (x′

v, y′

v < H and 0 ≤ y′

v, y′
v).
v < H.

The new position (x′
as follows:

v, y′

v) is determined by the vehicle’s current position (xv, yv) and orientation θv

(x′

v, y′

v) =





(xv, yv − 1)
(xv, yv + 1)
(xv − 1, yv)
(xv + 1, yv)

if θv = up
if θv = down
if θv = left
if θv = right

(3)

Interventions The intervention process follows a specific sequence: first, a step in the environment
dynamics is executed; then, an intervention is applied; finally, a snapshot of the resulting state is
captured. Interventions can modify traffic light states, alter obstacle positions, or move a vehicle
forward. Spatial interventions on obstacles and vehicles are constrained to single-cell displacements;
for obstacles, the direction is stochastic, while for vehicles, it is deterministically forward. Vehicle
intervention is further constrained by the absence of obstacles or traffic lights in the target cell,
adherence to environment boundaries, and the corresponding traffic light displaying a red signal.
A no-operation (NOOP) intervention is also permissible. This tripartite sequence—environmental
progression, intervention, and state documentation—constitutes a complete intervention cycle. These
interventions correspond to regime variables Rt, which are then represented using natural language.

Causal Variables The causal variables in the gridworld environment are the positions of the
vehicles (xv, yv), the positions of the obstacles (xo, yo), and the states of the traffic lights stl.

16

810
811

812
813
814
815
816
817

818
819
820
821
822
823

824
825
826
827
828

829
830
831
832
833
834

835
836
837
838
839
840

841
842
843
844
845
846

847
848
849
850
851

852
853
854
855
856
857

858
859
860
861
862
863

Under review as a conference paper at ICLR 2025

864
865

866
867
868
869
870
871

872
873
874
875
876
877

878
879
880
881
882

883
884
885
886
887
888

889
890
891
892
893
894

895
896
897
898
899
900

901
902
903
904
905

906
907
908
909
910
911

912
913
914
915
916
917

B ITHOR KITCHEN ENVIRONMENT - EMBODIED AI

The iTHOR (Kolve et al., 2017) kitchen environment is based on the FloorPlan10 dataset, featuring
a static 3D-rendered kitchen. The robot’s position remains fixed in front of the kitchen counter.
The environment consists of C underlying causal variables that interact based on the actions taken
by the agent. The environment contains three types of entities: movable objects m ∈ M , fixed
interactive objects f ∈ F , and receptacles r ∈ R. Movable objects include a plate with a potato and
an egg. Fixed interactive objects comprise a microwave, stoves, cabinet, and toaster. Receptacles
include the counter, microwave (when open), and pan (for the egg). Each object has a state so ∈ So,
where So is the set of possible states for object o. For binary state objects (e.g., microwave, cabinet),
So = open, closed or active, inactive. For movable objects, So includes their position (xm, ym, zm)
in the 3D space and a binary pickup state. The set of possible actions A includes:

• ToggleObject(o): For o ∈ {microwave, stoves, toaster}
• OpenObject(o): For o ∈ {microwave, cabinet}
• PickupObject(m): For m ∈ M
• PutObject(m, r): For m ∈ M, r ∈ R
• MoveObject(m): For m ∈ M
• NoOp: No action performed

The availability of actions depends on the current state of objects. For example:

ToggleObject(microwave) is valid iff smicrowave = closed
OpenObject(microwave) is valid iff smicrowave = inactive
(5)
The regime variable Rt ∈ [0, 1]2 represents the normalized click-location on the image to select the
object for interaction. Let Io be the set of pixels belonging to object o in the current frame. Then:

(4)

Rt =

1
H × W

· (x, y), where (x, y) ∼ Uniform(Io)

(6)

where H and W are the height and width of the frame respectively. The causal variables C =
C1, ..., CC in this environment correspond to the states and positions of objects. Binary state variables
(e.g., Cabinet-Open, Microwave-Active) take values in 0, 1, while position variables (e.g., Egg-Pos-x)
take continuous values in [0, 1], normalized to the environment’s dimensions. Observations are
generated as high-resolution images X t ∈ R512×512×3, then downsampled to X ′t ∈ R256×256×3
using bilinear interpolation.

C TEXT-BASED ACTION REPRESENTATION GENERATION

C.1 GRIDWORLD ENVIRONMENT

For the GridWorld environment, we implement a probabilistic context-free grammar (PCFG). The
PCFG includes:

• A set of adjectives Ao for each object type o ∈ O = traffic light, vehicle, obstacle
• A set of action modifiers M
• A set of action verbs Va for each action type a ∈ A = move, turn, change state

Let C : R3 → Σc be a function mapping RGB values to a finite set of color names Σc. For each
object o with RGB value ro, we compute its color name as co = C(ro). The generation process for
an action a on object o can be formalized as:

D(a, o) = m · va · the · adjo · co · o
(7)
where m ∼ P (M ), va ∈ Va, adjo ∼ P (Ao), and P (·) denotes the probability distribution defined
by the PCFG. Example: Consider an action to move a blue car to the right. Let ro = (0, 0, 255),
C(ro) = “blue”, m = “skillfully”, va = “moved”, and adjo = “sleek”. The generated description
would be:

D(move right, car) = “You skillfully moved the sleek, blue car to the right.”

(8)

17

Under review as a conference paper at ICLR 2025

C.2

ITHOR ENVIRONMENT

For the iTHOR environment, we define a mapping function f : A × O → Σ, where A is the set
of possible actions, O is the set of objects, and Σ is the set of all possible strings over the alphabet.
Let Ta : A → V be a function that maps actions to verb phrases, and To : O → Σ∗ be a function
that maps objects to descriptive phrases. The generation process for an action a on object o can be
expressed as:

f (a, o) = You · Ta(a) · To(o)
Example: For the action of opening a microwave, let a = OpenObject and o = Microwave. As-
sume Ta(OpenObject) = “adjusted” and To(Microwave) = “the microwave’s door”. The generated
description would be:

(9)

f (OpenObject, Microwave) = “You adjusted the microwave’s door.”

(10)

C.3 TOKENIZATION AND INTEGRATION

Let τ : Σ∗ → Nk be a tokenization function that maps a string to a sequence of k token indices. For
a generated description d, we compute its tokenized representation as:

t = τ (d)

(11)

The tokenized representations are then padded or truncated to a fixed length l, resulting in the final
representation t′ ∈ Nl. For a trajectory of actions a1, ..., an on objects o1, ..., on, we generate a
sequence of tokenized descriptions t′

1, ..., t′

n, where:

t′
i = pad(τ (D(ai, oi)), l) for GridWorld

t′
i = pad(τ (f (ai, oi)), l) for iTHOR

(12)

(13)

D DATA GENERATION AND PREPARATION

For each environment, we generated multiple datasets as shown in Table 5.

Table 5: Dataset specifications for each environment

Dataset

Size

Description

Training
Validation
Test
ICL
N -step evaluation

10000 trajectories of 100 steps
Used for model training
1000 episodes of 100 steps
Used for model validation
1000 episodes of 100 steps
Used for final evaluation
100 episodes of 100 steps
Used for in-context learning
100 episodes of 100 steps each, Used for N -step experiments
for each N value

D.1 DATA GENERATION PROCESS

The data generation process for both environments follows these steps:

1. Initialize the environment state randomly, ensuring a valid starting configuration.

2. For each step in the trajectory:

(a) In the Gridworld environment, apply the dynamic update rules (e.g., moving vehicles

if facing a green light).

(b) Select a random valid action from the set of possible actions for the current state.
(c) Apply the selected action to the environment.
(d) Record the current state, action taken, and resulting next state.

3. Repeat step 2 for the desired number of steps (100 in our case).

18

918
919

920
921
922
923
924
925

926
927
928
929
930
931

932
933
934
935
936

937
938
939
940
941
942

943
944
945
946
947
948

949
950
951
952
953
954

955
956
957
958
959

960
961
962
963
964
965

966
967
968
969
970
971

Under review as a conference paper at ICLR 2025

972
973

974
975
976
977
978
979

980
981
982
983
984
985

986
987
988
989
990

991
992
993
994
995
996

997
998
999
1000
1001
1002

1003
1004
1005
1006
1007
1008

1009
1010
1011
1012
1013

1014
1015
1016
1017
1018
1019

1020
1021
1022
1023
1024
1025

For the Gridworld environment, valid actions include toggling traffic lights and performing no action.
The dynamic nature of this environment means that even when no action is taken, the state may
change due to vehicle movements.

For the iTHOR environment, valid actions depend on the current state and may include toggling
objects (e.g., microwave, stoves), opening objects (e.g., cabinet), picking up or putting down movable
objects, and performing no action.

For N -step experiments, we generate multiple datasets, each corresponding to a different value of N :

• Gridworld: We create separate datasets for N ∈ {2, 4, 6, 8}.

• iTHOR: We create separate datasets for N ∈ {2, 4}.

Each N -step dataset consists of 100 episodes, where each episode is created by splicing together N
consecutive steps from the evaluation datasets. This approach provides sequences of varying temporal
lengths for our experiments.

D.2

IN-CONTEXT LEARNING EXAMPLES

For Gridworld, we maintain a pool of 10 ICL examples, each consisting of a 3-tuple (ini-
tial_state_causal_variables, actions, end_state_causal_variables). For each iteration during training
or evaluation, we randomly sample two examples from this pool to provide context for the model.
This process is similar to the one employed in RAP (Hao et al., 2023).

For iTHOR, we craft fixed few-shot examples to ensure comprehensive coverage of state-action pairs.
The examples are designed to demonstrate various object interactions and their outcomes. For 2-step
experiments, we use 7 examples covering every state-action pair at least once. For 4-step experiments,
we use 9 examples covering at least 2 of each state-action pair. This approach ensures that the model
has exposure to a wide range of possible interactions within the environment.

D.3 SELF-EVALUATION REWARDS

Following RAP (Hao et al., 2023), for the self-evaluation rewards in planning tasks, we generate
samples by splicing 1-step trajectories. We select the actual action taken in the environment for
“good” evaluations, providing a positive example of a correct action. For “bad” evaluations, we select
a random action different from the one actually taken, providing a negative example.

E CRL MODEL TRAINING

This section details the training process for the Causal Representation Learning (CRL) models
used in our experiments. The CRL models are trained using triplets of (state_image, text action,
next_state_image) following the process described in Section 4.

E.1 AUTOENCODER

The autoencoder is trained from scratch using 10 times more samples than the main dataset to ensure
a robust representation. This approach is justified by the relative ease of obtaining unlabeled, random
samples from an environment. In scenarios where this is not feasible, transfer learning from a
pretrained image representation model can be employed by adding a learnable linear projection to the
required dimensions and training with the original dataset size.

For the Gridworld environment, we implement an autoencoder with 40 latent dimensions and 64
hidden channels. Both the encoder and decoder consist of 2 residual blocks with SiLU activation
functions. We incorporate the CoordConv operator (Liu et al., 2018) to better capture coordinate
information from images. For the iTHOR environment, we employ the autoencoder architecture from
BISCUIT (Lippe et al., 2023).

19

Under review as a conference paper at ICLR 2025

1026
1027

1028
1029
1030
1031
1032
1033

1034
1035
1036
1037
1038
1039

1040
1041
1042
1043
1044

1045
1046
1047
1048
1049
1050

1051
1052
1053
1054
1055
1056

1057
1058
1059
1060
1061
1062

1063
1064
1065
1066
1067

1068
1069
1070
1071
1072
1073

1074
1075
1076
1077
1078
1079

E.2 NORMALIZING FLOW AND TRANSITION MODEL

For both the normalizing flow and transition model, we use the same architectures and hyperparam-
eters as in BISCUIT (Lippe et al., 2023) as it has demonstrated strong performance in identifying
causal variables from high-dimensional observations.

E.3 TEXT ENCODER

The text encoder for the Gridworld environment is based on a pretrained Sentence Transformer
(Reimers & Gurevych, 2019), specifically the all-MiniLM-L6-v2 model3, augmented with a 2-layer
MLP head with 64 hidden dimensions. For iTHOR, we use a pretrained SigLIP model (Zhai et al.,
2023)4 with a similar 2-layer MLP head. In both cases, the pretrained encoders remain frozen during
training, with only the MLP head being updated.

E.4 TRAINING PARAMETERS

Key training parameters for each environment are as follows: For Gridworld, we use a learning rate
of 3 × 10−3 for the main model and 3 × 10−3 for the text MLP, batch size of 384, and train for 300
epochs. For iTHOR, we use a learning rate of 1 × 10−3 for the main model and 3 × 10−3 for the text
MLP, batch size of 64, and train for 100 epochs. Both environments employ a warmup period of 100
steps and a sequence length of 2 for training.

F MODEL SELECTION

This section details our model selection procedure for the different components of our framework.

F.1 MODEL COMPONENTS

For the text encoder, we performed 5-fold cross-validation to select the optimal hyperparameters for
the MLP head architecture and training parameters. The search parameters for the planning algorithm
were optimized using Bayesian optimization with 15 trials.

F.2 SENSITIVITY

Our experiments indicated that the framework’s performance is relatively robust to variations in
the model training hyperparameters. The causal encoder and text encoder components showed
stable performance across different configurations. However, we observed higher sensitivity to the
exploration weight parameter w in the search algorithm due to the interaction between exploration-
exploitation trade-off and reward scaling.

G CAUSAL MAPPER

The causal mapper mθ extracts interpretable causal variables from the learned disentangled represen-
tations.

This process allows for a non-injective mapping from latent dimensions to causal variables. For
instance, if we have a causal variable “cabinet_state”, the first stage might learn that latents 1, 5, and
7 are the most predictive for this variable. In the second stage, a specific predictor would learn to
map from these dimensions to either 0 (closed) or 1 (open).

The causal mapper mθ is implemented in two stages:

G.1 TARGET ASSIGNMENT

This stage uses a single MLP fassign to predict all causal variables from each latent dimension
independently:

3https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
4https://huggingface.co/timm/ViT-B-16-SigLIP

20

Under review as a conference paper at ICLR 2025

ˆC = fassign(z ⊙ M, M )

(14)

where M ∈ {0, 1}L is a mask and ⊙ is element-wise multiplication. M is set to the identity.

For each latent dimension i, we create a mask Mi where only the i-th element is 1 and the rest are 0.
We then batch these masks along with the corresponding masked latent vectors:







z ⊙ M1
z ⊙ M2
...
z ⊙ ML







(15)







,







M1
M2
...
ML

This batched input is fed into fassign, which outputs predictions for all causal variables for each
masked input. The output shape is [L, C], where L is the number of latent dimensions and C is the
number of causal variables.

We then compute the correlation between these predictions and the ground truth causal variables.
This allows us to identify which latent dimensions are most predictive of each causal variable. We
apply a correlation threshold (in our experiments we use 0.1) to determine which latent dimensions
are relevant for each causal variable to determine each M ′
j.

G.2 CAUSAL PREDICTION

Individual MLPs fcausal,j are trained for each causal variable j, using only the relevant latent dimen-
sions identified in stage 1:

ˆCj = fcausal,j(z ⊙ M ′
j)

(16)

where M ′

j is the mask for causal variable j.

The output layer of each fcausal,j is adjusted based on the a priori known type of the causal variable
(categorical, numerical, angle).

H STATE DESCRIPTION GENERATOR

The state description generator s is responsible for converting the causal variables into a human-
readable natural language description of the current state. This process can be implemented in various
ways:

H.1 STOCHASTIC AND DETERMINISTIC IMPLEMENTATIONS

The generator can operate either stochastically or deterministically, depending on the application’s
needs:

1. Stochastic: This approach uses a language model with a temperature greater than 0, which
allows for a variety of possible descriptions for the same state. This variability can be useful
in scenarios where diverse language outputs are desired.

2. Deterministic: This method involves either setting the temperature of a language model
to 0, ensuring consistent outputs, or using a rule-based system that directly maps causal
variables to fixed phrases or sentences.

H.2 EXAMPLE OF STATE DESCRIPTION GENERATION

For instance, given a dictionary of causal variables as follows:

21

1080
1081
1082

1083
1084
1085
1086
1087
1088

1089
1090
1091
1092
1093
1094

1095
1096
1097
1098
1099
1100

1101
1102
1103
1104
1105

1106
1107
1108
1109
1110
1111

1112
1113
1114
1115
1116
1117

1118
1119
1120
1121
1122
1123

1124
1125
1126
1127
1128

1129
1130
1131
1132
1133

Under review as a conference paper at ICLR 2025

{

}

"cabinet_state": 1,
"light_color": 0,
"door_angle": 45

The state description generator might produce a sentence like: "The cabinet is open. The traffic light
is showing a red signal. The door is partially open at a 45-degree angle."

H.3 CHOICE OF APPROACH

The choice between a stochastic and deterministic approach depends on the specific requirements
of the task and the desired level of variability in the generated descriptions. For simplicity and
consistency, in our experiments, we have opted for a rule-based deterministic state descriptor.

While rule-based generation is suitable for environments with reasonably sized state spaces, more
complex environments may benefit from learned approaches. A fine-tuned sequence-to-sequence
model or instruction-tuned LLM could generate natural descriptions from causal variables while
maintaining consistency. The key requirement is that the mapping from causal variables to descrip-
tions remains reliable and interpretable, allowing the planning agent to reason effectively about state
transitions.

The modular nature of our framework allows for easy substitution of the state description generator.
This flexibility ensures that as environments become more complex, the description generation can
be adapted accordingly while maintaining the benefits of our causally-aware planning approach.

I ENABLING THE BASELINE TO PROCESS IMAGE STATES

To enable LLaMA to process the environment states, we implement a conversion of visual states
to natural language descriptions using the ground truth causal variables. This process ensures fair
comparison with our causal world model while maintaining the LLM’s ability to reason about the
environment.

For each initial state, we extract the ground truth causal variables and use the same rule-based
state description generator employed in our causal world model to convert these variables into
2,
natural language. For example, in GridWorld, a state with causal variables blue_car_x:
blue_car_y:
green would be converted to “The blue car is
at position (2,3). The cyan traffic light is showing green.”.

3, cyan_light_state:

The baseline LLM then uses this initial state description to reason about subsequent states and actions,
relying on its world model capabilities to predict state transitions. This approach ensures that the
baseline has access to the initial information as our causal world model, with the key difference being
that our model learns the causal structure while the baseline relies on its pretrained knowledge for
state transition predictions.

J CAUSAL MAPPER ANALYSIS

We present a statistical analysis framework to evaluate the performance of our causal mapper in the
GridWorld environment.

J.1 ANALYSIS METHODOLOGY

Our evaluation framework consists of two core components:

1. Overall Performance Analysis: We track mean absolute error (MAE) across all dimensions
against training set size. Standard deviation bands are computed from three independent
training runs to illustrate the variance in performance across different training instances.

22

1134
1135

1136
1137
1138
1139
1140
1141

1142
1143
1144
1145
1146
1147

1148
1149
1150
1151
1152

1153
1154
1155
1156
1157
1158

1159
1160
1161
1162
1163
1164

1165
1166
1167
1168
1169
1170

1171
1172
1173
1174
1175

1176
1177
1178
1179
1180
1181

1182
1183
1184
1185
1186
1187

Under review as a conference paper at ICLR 2025

2. Dimension-wise Evolution Analysis: We analyze how prediction accuracy for each causal
variable evolves with training size using heatmaps, with darker colors indicating better
performance.

Statistical significance is assessed using the criterion that standard deviation should be less than half
the mean value, indicating reliable performance measurements.

(a) Model performance (MAE) vs training size for the
causal mapper in GridWorld. Error bands represent one
standard deviation across three independent runs.

(b) Dimension-wise performance analysis showing the
evolution of prediction accuracy across dimensions at
different training sizes. Darker green colors indicate
better performance.

Figure 3: Performance analysis of the causal mapper showing both overall error metrics and
dimension-wise evolution.

J.2 RESULTS

Our analysis reveals strong performance of the causal mapper approach in terms of data efficiency and
prediction accuracy. The causal mapper achieves adequate accuracy (MAE < 0.05) with approximately
1200 labeled examples, demonstrating effective learning from disentangled representations.

The dimension-wise evolution analysis reveals distinct learning patterns across different types of
causal variables. The causal mapper exhibits rapid early learning for traffic light states, achieving
high accuracy with minimal data. For positional variables, we observe more gradual but consistent
improvement as training data increases. This pattern suggests that binary state variables (like traffic
light states) are easier to learn than continuous positional variables, which require understanding
more complex spatial relationships.

The performance analysis shows consistent improvement across all dimensions as training size
increases, with particularly strong performance in predicting traffic light states even in low-data
regimes. The small standard deviation bands indicate stable learning across different training runs,
suggesting robust performance regardless of initialization conditions.

K EVALUATION METHODOLOGY

Given the stochastic nature of both the Gridworld and iTHOR environments, we have implemented
specific adjustments to our evaluation methodology. These adjustments ensure that our perfor-
mance metrics accurately reflect the models’ understanding of the underlying causal structure while
accounting for inherent randomness in the environments.

K.1 GRIDWORLD ENVIRONMENT

In the Gridworld environment, we make the following adjustment:

23

1188
1189

1190
1191
1192
1193
1194
1195

1196
1197
1198
1199
1200
1201

1202
1203
1204
1205
1206

1207
1208
1209
1210
1211
1212

1213
1214
1215
1216
1217
1218

1219
1220
1221
1222
1223
1224

1225
1226
1227
1228
1229

1230
1231
1232
1233
1234
1235

1236
1237
1238
1239
1240
1241

102103Training Set Size0.020.040.060.080.100.120.140.16Mean Absolute ErrorOverall PerformanceCausal Mapper100% Training Data (n=16000.0)80961121281441602403204004805606407208008809601040112011841200128013601440152016001920238424003200528016000Training Sizeobstacle (255, 165, 0) position xobstacle (255, 165, 0) position ytrafficlight (0, 255, 255) statetrafficlight (100, 100, 0) statetrafficlight (192, 192, 192) statevehicle (0, 0, 255) position yvehicle (192, 192, 192) position yvehicle (255, 0, 0) position yDimensionCausal Mapper Dimension Evolution0.000.050.100.150.200.250.30Mean Absolute ErrorUnder review as a conference paper at ICLR 2025

1242
1243

1244
1245
1246
1247
1248
1249

1250
1251
1252
1253
1254
1255

1256
1257
1258
1259
1260

1261
1262
1263
1264
1265
1266

1267
1268
1269
1270
1271
1272

1273
1274
1275
1276
1277
1278

1279
1280
1281
1282
1283

1284
1285
1286
1287
1288
1289

1290
1291
1292
1293
1294
1295

• Boulder Position Exclusion: We exclude the boulder’s position from the final state evalua-
tion. This is because the boulder’s movement is stochastic and not determined by the causal
structure we aim to learn and evaluate other than the fact that it was moved or not.

Rationale: The boulder’s position can vary due to random factors not captured in our causal model.
By excluding it from our evaluation, we focus on the aspects of the environment that are causally
determined by the actions and states we’re modeling.

K.2

ITHOR ENVIRONMENT

For the iTHOR environment, we implement a more nuanced approach:

• Coordinate Categorization: We categorize the x, y, and z coordinates for objects with

stochastic movements into discrete position categories.

• Category-based Evaluation: Instead of comparing exact coordinates, we check whether

objects end up in the correct category of positions after an action.

Rationale: In iTHOR, object movements can have slight variations due to a) inherent stochasticity
in movements, and b) physics simulations, even when the same action is applied. By categorizing
positions, we can evaluate whether the model correctly predicts the general outcome of an action
(e.g., “on the counter” vs. “in the microwave”) without being overly sensitive to minor coordinate
differences.

K.3

IMPLEMENTATION DETAILS

For both environments, we implement these adjustments as follows:

1. State Representation and Prediction: We maintain the full state representation, including

all object positions and attributes, for both the actual and predicted states.

2. Dynamic Evaluation: During the comparison of predicted states to ground truth states, we

dynamically apply our adjustment rules:

• For Gridworld, we dynamically ignore the boulder’s position when comparing states.
• For iTHOR, we dynamically categorize the exact x, y, z coordinates into position
categories (e.g., “on the counter”, “in the microwave”) and compare these categories
instead of the exact coordinates.

3. Accuracy Calculation: We calculate accuracy based on the match between predicted and
actual states after applying these dynamic adjustments during the comparison process.

L COMPUTATIONAL OVERHEAD ANALYSIS

We performed detailed benchmarks comparing single-step predictions between the LLM-based world
model and our CRL world model. The analysis was conducted on 5 GridWorld samples, with 10 runs
per sample after warmup, using a single NVIDIA A100-40GB GPU.

Our CRL world model consists of three main components: an autoencoder (4.5M parameters), a
normalizing flow (2.9M parameters), and a transition prior (28.7M parameters), totaling 36.1M
parameters. For comparison, we used LLaMA 3 8B quantized to 6 bits via ExLlamaV2 as our
baseline LLM world model.

The benchmarks revealed that the CRL world model achieves an average inference time of 27ms,
compared to 2.2s for the LLM world model—representing an approximately 82x speedup. This
computational difference has significant implications for planning tasks. For example, a 10-branch
tree search would take approximately 22 seconds with LLM calls versus just 0.27 seconds with the
CRL world model.

This substantial performance difference becomes particularly important in scenarios requiring multiple
forward simulations or when real-time planning is necessary. The efficiency of our CRL world model
enables more extensive tree searches and faster iteration during planning, while maintaining high
prediction accuracy as demonstrated in our main experimental results.

24

Under review as a conference paper at ICLR 2025

M CAUSAL WORLD MODEL ALGORITHM

This section presents the formal algorithm for sampling from/performing inference with the Causal
World Model. The algorithm takes as input the trained model components and an initial state, and
produces a sequence of latent states and their corresponding natural language descriptions.

Algorithm 1 Inference with the Causal World Model
Require: Observation space X , latent space Z, action description space L, action encoding space
A; observation encoder eψ, normalizing flow fϕ, action encoder Le, transition model pω, causal
mapper mθ, state description generator s; initial observation X0 ∈ X ; action descriptions
{Lt}T −1

▷ Causal encoding of observation

t=0 ∈ LT
1: function E(X ∈ X )
E ← fϕ(eψ(X))
2:
return E
3:
4: end function
5: function ENCODEACTION(L ∈ L)
a ← Le(L)
6:
return a
7:
8: end function
9: function G(z ∈ Z)
C ← mθ(z)
10:
ℓ ← s(C)
11:
return ℓ
12:
13: end function
14: function SAMPLENEXTSTATE(zt ∈ Z, at ∈ A)
15:
16:
17:
18:
19: end function
20: function INFERENCETRAJECTORY(X0 ∈ X , {Lt}T −1
t=0 )
21:
22:
23:
24:
25:
26:
27:
28:
29:
30: end function

at ← EncodeAction(Lt)
(zt+1, ℓt, ℓt+1) ← SAMPLENEXTSTATE(zt, at)
yield (zt+1, ℓt+1)
zt ← zt+1

zt+1 ∼ pω(zt+1 | zt, at)
ℓt ← G(zt)
ℓt+1 ← G(zt+1)
return (zt+1, ℓt, ℓt+1)

z0 ← E(X0)
ℓ0 ← G(z0)
yield (z0, ℓ0)
for t = 0 to T − 1 do

end for

▷ Encode action description into action latent space

▷ Map latent state to causal variables
▷ Generate natural language state description

▷ Predict next latent state
▷ Generate current state description
▷ Generate next state description

▷ Initialize latent state
▷ Generate initial state description

▷ Encode action description

▷ Update latent state for next iteration

1296
1297

1298
1299
1300
1301
1302
1303

1304
1305
1306
1307
1308
1309

1310
1311
1312
1313
1314

1315
1316
1317
1318
1319
1320

1321
1322
1323
1324
1325
1326

1327
1328
1329
1330
1331
1332

1333
1334
1335
1336
1337

1338
1339
1340
1341
1342
1343

1344
1345
1346
1347
1348
1349

25

Under review as a conference paper at ICLR 2025

N MODIFIED MCTS PLANNING ALGORITHM

We adapt the Reasoning via Planning (RAP)(Hao et al., 2023) Monte Carlo Tree Search (MCTS)
algorithm (Kocsis & Szepesvári, 2006; Coulom, 2006) for our causally-aware planning framework.
Our modifications primarily focus on integrating the causal world model and leveraging its capabilities.
Algorithm 2 presents our modified version of the MCTS algorithm.

Algorithm 2 Causally-Aware MCTS
Require: Initial image X0, causal world model (Algorithm 1), LLM agent, depth limit L, number
ICL, self-evaluation ICL samples

of roll-outs N , exploration weight w, intuition ICL samples Dint
Dself
ICL

1: Initialize memory of actions A : Z (cid:55)→ L, children c : Z × L (cid:55)→ Z and rewards r : Z × L (cid:55)→ R
2: Initialize the state-action value function Q : Z × L (cid:55)→ R and visit counter N : Z (cid:55)→ N
3: z0 ← E(X0), ℓ0 ← G(z0)
4: for n ← 0, . . . , N − 1 do
5:
6:
7:

t ← 0, zt ← z0, ℓt ← ℓ0
while N (zt) > 0 do

▷ Initialize root node

▷ Selection

(cid:113) ln N (zt)
N (c(zt,a))

(cid:105)

N (zt) ← N (zt) + 1
at ← arg maxa∈A(zt)
rt ← r(zt, at), zt+1 ← c(zt, at)
t ← t + 1, zt ← zt+1, ℓt ← G(zt)

(cid:104)
Q(zt, a) + w

end while
while zt is not a terminal state ∧ t ≤ L do

At ← GetValidActions(ℓt)
for a ∈ At do

zt+1, _, ℓt+1 ← SAMPLENEXTSTATE(zt, Le(a))
rintuition ← − log pLLM(a | ℓt, Dint
rself-eval ← − log pLLM(“good” | ℓt, a, Dself
ICL)
r(zt, a) ← rintuition + rself-eval
Update A(zt) ← A(zt) ∪ {a}, c(zt, a) ← zt+1

ICL)

end for
at+1 ← arg maxa∈A(zt) r(zt, a)
rt ← r(zt, at+1), zt+1 ← c(zt, at+1)
t ← t + 1, zt ← zt+1, ℓt ← G(zt)

end while
for t′ ← t, . . . , 0 do
Update Q(zt′

end for

, at′) with {rt′, rt′+1, . . . , rt}

8:

9:
10:
11:
12:
13:
14:
15:
16:
17:
18:
19:
20:
21:
22:
23:
24:
25:
26:
27:
28: end for

▷ Expansion
▷ Use the Causal World Model

▷ Simulation

▷ Back propagation

The key modifications in our algorithm compared to the original RAP MCTS are:

1. State Representation: We use disentangled causal latent representations z for states, starting

from an encoded initial image X0 (line 3).

2. Causal World Model Integration: We employ our trained causal world model (Algorithm

1) to predict the next state and generate state descriptions (line 15).

These modifications allow our MCTS algorithm to leverage the causal understanding provided by
the causal world model, while also incorporating the strengths of the LLM agent for action selection
and evaluation. The use of disentangled latent representations z allows for efficient and robust state
transitions, while the natural language descriptions ℓ enable interaction with the LLM agent.

While our current implementation uses a predefined set of valid actions, the framework could
potentially be extended to sample actions directly from the LLM for open-ended domains where the
action space is not easily enumerable.

26

1350
1351

1352
1353
1354
1355
1356
1357

1358
1359
1360
1361
1362
1363

1364
1365
1366
1367
1368

1369
1370
1371
1372
1373
1374

1375
1376
1377
1378
1379
1380

1381
1382
1383
1384
1385
1386

1387
1388
1389
1390
1391

1392
1393
1394
1395
1396
1397

1398
1399
1400
1401
1402
1403

