Language Agents Meet Causality -- Bridging LLMs and Causal World Models
Authors: John Gkountouras
Matthias Lindemann
Phillip Lippe
Efstratios Gavves
Ivan Titov
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Models
Causality
Causal Representation Learning
Language Agents
Planning
Primary area: applications to robotics, autonomy, planning
TL;DR: Improving LLM planning capabilities using learned causal representations

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.

Published as a conference paper at ICLR 2025
LANGUAGE AGENTS MEET CAUSALITY BRIDGING LLMS AND CAUSAL WORLD MODELS
John Gkountouras', Matthias Lindemann', Phillip Lippe', Efstratios Gavves and Ivan Titov
Institute for Logic, Language and Computation (II.CC), University of Amsterdam Institute for Language, Cognition and Computation (IICC), University of Edinburgh QQUA Lab, University of Amsterdam
Arrchimeess/Athene RC, Greece
i. www.untouussuuaa.n.uu. m. m. www.dddannnnnmss.... ed. ac, uk,
p p.lippeeuva.nl, egavveseuva.ni nl, ititovêinf edd ac
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 specific environment. In contrast, causal representation learning (CRL) focuses identifying the underlying causal structure within a given environment. We propose framework that integrates CRLs with LLMs 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.
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 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-specifie, downstream performance (Ahn)
et [0.22)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 its training data (Zečević 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 identify the underlying causal structure of data (Schölkopf et al.) 2021). By separating and identifying latent causal factors, CRL enables models 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 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 Al systems, Project page: http:://ommgguugghhhhoooLLMMMMW

---

Published as a conference paper at ICLR 2025
Causel Nodal
Gausal Model
Crucal
Text
Teet Enecate
E
Tiet
Encodar
plas dows
Figure 1: Overview of a single rollout in the proposed planning pipeline. The causal encoder, implemented using CRL model, maps the high-dimensional state representation (image) to its fundamental constituents-uue 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.
including LLMs. Although CRL does not directly address all LLM limitations, 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 the L.M 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 flexible and intuitive way describe actions, making them more suitable for generalist agents operating in diverse environments. Moreover, annotating frame sequences with natural language descriptions often casier than exhaustively enumerating every possible action in an environment, which can be intractable for complex domains.
We consider a setting with interleaved sequential observations 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 series of images showing the robot's workspace, paired with descriptions like "The gripper shifted slightly to the right." "The object was grasped and placed the worktop." We assume no prior knowledge the causal factors or the causal mechanisms between them. The agent 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 BISCUIT (Lippe et al. 2023). CRL framework, to create a flexible causally-awan 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 current frontier of causal representation learning. As more powerful CRL methods become available, they can be integrated into framework, scaling it to more complex, realistic scenarios.
case. While these environments are still relatively simple, they represent the
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
Causal Representation Learning Causal representation learning aims to identify the underlying causal variables and their relations from high-dimensional observations (Schölkopl et al. 202I). In the most general setting, the latent causal variables may not be uniquely identifiable (Locatello
et al. 2019a; Hyvärinen & Pajunen| 1999). Many approaches rely 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 al., 2023 Kivva al. [2022] Lachapelle et al.
2023). ) sparse graphical structures (Khemakhem et al. [2000] Liu et al. [022||0224 Lachapelle &
Lacoste-Julien 2022; Lachapelle et al. 2024). having muluple views (Xu et al. 2024] Yao et al.
2024a| von Kügelgen et 20.] Brehmer et al. 2022, Locatello et al. (2020), or supplementary
supervision labels (Yang et al. [2200] Komanduri et al., [0022] 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),
CRL framework that learns causal representations with realistic agent-focused assumptions, requiring only 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) (018)) Recent work has
focused on object-centric world models (Greff et al., 2017 Steenkiste et al. [0188] Watters 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 causal world model relying only on images and textual annotations but capable of reasoning about actions across state transitions, while also being able to be interacted with by 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. (023). Various methodologies have been developed to make use of LLMs for agent planning. These include task decomposition for breaking 2023 Shen et al. (2024), multi-plan selection Yao al.
complex tasks into subtasks (Wei et al. 2022 Yao et al.
for generating and choosing optimal plans
2024b Wang et al. 20022, external module- reflection and refinement via self-evaluation and
aided planning (Liu et al.] 2023a Guan et al.
2023).
improvement (Shinn al. 2024 Gou al.
2024
Madaan et al., 2024). and -aggneeed
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-follomin
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) & Everiti 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).
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.
Except for few labels needed to map from the latent representation human-interpretahle language.
3

---

Published as a conference paper at ICLR 2025
3.1 CAUSAL REPRESENTATION LEARNING IN TEMPORAL CAUSAL GRAPHS
CRL aims to uncover the latent causal variables and the underlying causal structure. In temporal settings, we consider sequences of high-dimensional observations (X')=>> where X RP, and actions {R where R' E R Actions R' can represent, for example, the coordinates of the locations where the interactions occurred (Lippe al. 2023). The true causal variables {C*) =t=0 where C E R* are unobserved. Furthermore, deterministic observation model is assumed, often represented as X' g(C), where g R.K RDD is an injective function mapping causal variables to observations.
Instead of directly modeling causal variables, CRL relies on latent state representations. estimates function ʃ RD > M that maps observations X' to latent representations z' [(X*). The goal is to ensure that each dimension zद of z' corresponds to causal variable C₁ in C' up to a transformation decided by the identifiability class of the causal model. Specifically, it aims to achieve this disentanglement using only {X' 17 and 1/
3.2 GENERATIVE MODEL
The temporal CRL framework often modeled as a generative process that describes how observa- tions are produced from underlying latent state representations and actions. At each time step i, the state z' evolve according to a transition model influenced by actions R', and generate observations X Assuming first-order Markov process, the conditional likelihood of the observed data {X') { given actions (R' Ji=1 expressed as
T
-f p(22) II Pz x'-1 RR) (
2²) dz,
=
where p(z") is the prior distribution over the state. The transition model term (z models the state dynamics, capturing how the states evolve over time and how intervening actions influence them. The observation model Py (X' zz) describes how the states generate the observations, which in our case will be done with the deterministic function g.
( 1 R²)
The marginalization Eq. T renders the objective intractable. A standard approach to address this is to optimize the corresponding Evidence Lower Bound (ELBO) by assuming Gaussian distribution for the transition dynamies 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. ① 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 R on x' is mediated by latent binary mask, and (2) these interaction patterns vary over time. The first assumption is enforced by using structured model family to model the transition Pw, (z. x² - Rt). We incorporate this component from BISCUIT in approach. These assumptions together ensure that causal variables 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 the original paper (Lippe et al. 2023).
BUILDING CAUSAL WORLD MODEL FROM CAUSAL REPRESENTATIONS
To integrate the CRL model with LLMs, 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 's dynamics. with the CWM's '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/ineervntioons, as demonstrated by our experiments in Section 6
Since do normally have a priori information of the number of causal variables, M >> K.
4

---

Published as a conference paper at ICLR 2025
Input
Causal Encoder
Decoder
SECLIT
Model
F
Mapper
The open buuti light anoweng a posiiioed 10:15 75) The heaey touiter 0.77).
Mt
Model
State
You state th sear lgght
MIP
R
Target
Figure 2: Illustration of the first roll-out step with the Causal World Model. The image X and action description L 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 C' Finally, the state descriptor 8 generates a natural language description e of the next state. For subsequent steps, the model autoregress in the latent space using the previously predicted bypassing the autoencoder and normalizing flow, enabling efficient multi-step inference and planning.
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 2022
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 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-Basd Action Representations We replace the action encoding R the CRL frame- work with language-basee representation (L'),, where Le embeds natural language description L This is implemented using an encoder-only language model (Reimers & Gurevych 2019) with trainable head, replacing the original action encodings in the CRL framework's fransition model R' (L') (see also Section
Decoder The decoder comprises two parts: the causal mapper and the state description generator. The causal mapper my extracts causal variables C' from the learned disentangled representations z. It first identifies which latent dimensions 2; are most predictive for each causal variable Cj, then learns to perform the actual mapping. The state description generator maps the estimated causal variables to f, 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, 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 any environment.
To estimate the model, we use image pairs (X*) and corresponding action descriptions {L'} 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 X, into lower-dimensional latent representations E, e exx (X,), ). in which, however, the causal variables will still be entangled. Then, analogously to Eq. [ the conditional likelihood the encoded observations {E*) )
Estimation: Causal Encoder and Transition Model
given action descriptions {L' } is given by
T
((E' (L')) = f p(z)) II P z (--I,o(L')) (E 2²) dz,
where Pw is the transition model, structured as in BISCUIT in order to satisfy the identifiability guarantees, (E z²), is the observation model, and p(z") ) is the prior distribution over the
initial latents, assumed to be the standard Gaussian. The invertible mapping fa RM RM is
5

---

Published as a conference paper at ICLR 2025
normalizing flow (NF) that transforms the autoencoder's states (E*} } into new, structured latent space } while identifying and separating the causal variables. As the NF is invertible, fo and its Jacobian also yield the term P& (E* zz') in the generative model. Similar to BISCUIT, the ELBO is formulated and optimized using the reparameterization trick (Kingma & Welling 2013
The full causal encoder that maps observation X' to the causal state z is expressed as & == cg.. where Cy is the encoder part of the autoencoder, and So is the NF. While we assume perfect reconstruction capability for the autoencoder, common assumption in CRL (Kivva et al. [022 Lachapelle et al.] [0233 Brehmer el al., 2022 Lachapelle et al. 20244 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 (ReaINVP (Dinh et al. (016)), and transition model. However, we introduce an important modification the action representation. While BISCUIT used coordinate- based action encodings (as described in Section 5.1). our work incorporates language-basse action representations through (L') in the transition model to enable our model to process natural language action descriptions.
Estimation: Decoder We train the causal mapper my using a small set of annotated but not necessarily ordered images where the true causal variables C and their values known. The training pairs consist (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 21, 23. and 27 are most predictive for the "blue car x-position" (Co). and then train specific predictor for Co using only these relevant dimensions.
The state description generator s, typically rule-based system, maps the estimated causal variables to human iattannnnnterrrttallee outputs. For example, it might transform position and state variables into description like "The blue car is (2,3), the cyan traffic light green". The full decoder is expressed as G mg.
Inference Process During inference, the model sequentially processes new GridWorld images through these components:
z' = =((x') = o (x)) +1 ~ Pw (z (z zz Le,,e
g+ g(z = (s mz)
This process transforms raw input into interpretable state descriptions 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 X 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 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 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: dynamic 8 8 GridWorld and static 3D-rendered kitchen (A12-THÓR) (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 provided in Appendices and B

---

Published as a conference paper at ICLR 2025
For each environment, 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
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 position (2,3) 3) is transformed into higher-dimensional representation using high-frequency sinusoidal functions.
2. Text-based (TB): Natural language descriptions expanded using a PCFG (c.g.. "you toggled the bright cyan traffic light"), then encoded through an encoder-only text embedder.
3. Hybrid (HB): Concatenation of coordinate-bare 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 language- based interface suitable for 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 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 ae, and context c:
St411 PLM |te Oze c).
The bascline constructs prompt 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 ex12 format. We chose RAP+LIaMA3 the baseline for its simplicity and effectiveness, providing fair point of comparison to assess the benefits of integrating causal representation learning. This allows us isolate the impact of causal understanding 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 '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 standard CRL metric: Ri2 scores for the permutation ा that maximizes the diagonal of the R² matrix between learned latent variables and true causal variables. This approach accounts for the fact that 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 B
7

---

Published as a conference paper at ICLR 2025
Table 1: R² scores for action representations. CB: Coordinate-basee TB: Text-based, HB: Hybrid. 100% stands for 10t image states.
Subsample Percentage
Action Type
100%
0.3%
0.5%
0.7% %
1.0%
1.2%
1.5%
CB
0.392 +:0.000
0.366=00.000
0.424=0.001
0.457++0.004
0.472++0.004
0.548++0.022
0.987
TB
0.374=00000
0.302=00.000
0.399=0.000
0.470=0.012
0.495+0.0014
0.603+0.003
0.990
HB
0.392-10.000
0.433 0.000
0.461
0.490110.010 0.539-10.011
0.991
Table
N-step causal inference accuracies for the causal world model and the RAP Hao et al. world model across different environments and step lengths.
2023
iTHOR
Gril
Steps
1
2
4
8
Causal Model
0.824
0,680
0.630
0.954
0.922
0.829
0.797
0.758
RAP World Model
0.482
0.285
0.110
0.391
0.220
0.085
0.045
0.005
Table 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.
These findings support our hypothesis: action encodings including text are as effective as or superior in uncovering causal variables coordinate-base representations, particularly when data is scarce. We use the TB model trained using the entire dataset (10% 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 given sequence of actions rather than finding actions reach goal. The evaluation methodology is presented in Appendix K
Table prresentsss 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%) for 8-step inference, while the baseline's performance drops nearly to O. The performance in iTHOR, while lower than in GridWorld, still shows a substantial improvement over the bascline.
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 's lower performance GridWorld compared to iTHOR may be due to the lack of visual input, which is less natural for language models an artificial environment.
Table 3 provides provides detailed breakdown of the causal inference performance for specific actions and objects, based on the extended 1 -tspep dataset of 3000 samples. In iTHOR, the causal world model excels at ToggleObject and OpenOhject 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 performed intervention or not. Performance could be improved by injecting
to estimate whether inductive bias towards the continuous, diree-dimensionall nature the underlying variable. However, this requires task specialization within the model and chose keep the proposed framework task-
8

---

Published as a conference paper at ICLR 2025
Table Causal inference accuracies for action categories 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
0.957
0.466
Change Light State
0.986
0.300
OpenObject
0.926
0.339
No Action
0.985
0,456
NoOp
0.962
0.710
Move
0.928
0.408
PutObject
0.506
0.100
PickupObject 0.431
0.692
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 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 sequence of actions 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 language model agent, rather than using a single language model for both roles. We the same LLM as for the baseline planning agent (LLaMA 3).
The planning works as follows: The LLaMA 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 's quality and picks an action resulting in a new state. This cycle repeats, exploring multiple reasoning paths before converging
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
Actions In Gridworld, there are three actions to toggle traffic lights (one per light) and one to 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 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 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 correct, given the current state and few-shot demonstrations.
We avoid using problems that are easily divisible into subgoals subtasks. This choice ensures that our method remains applicable to wide range of problems, including those where intermediate progress toward the goal difficult quantify and/or may not correlate directly with overall success.
rewards to maintain generality and applicability to
6.3.1 PLANNING RESULTS AND DISCUSSION
Table 4 presents the planning results for both models across different environments and step lengths.

---

Published as a conference paper at ICLR 2025
Table 4: Planning results for the causal world model and RAP (Hao al. 2023 across different environments and step lengths. The best performing method for each metric is highlighted in green.
Causal World Model
RAP Baseline
Environment
Steps
Avg. Steps (Success)
Avg. Steps (Failure)
Avg. Steps (Success)
Avg. Steps (Failure) !
Success Rate
Success Rate
2
0.58
1.78
3,36
0.25
4.00
4.00
ITHOR
0.44
2.14
5.43
0.11
6.00
6.00
2
0.95
1.92
3.20
0.20
2.00
3.19
4
0.73
2.71
5.27
0.11
3.27
5.48
GridWorld
6
0.46
3.65
7.72
0.08
5,50
7.72
8
0.42
4.62
9.76
0.06
7,00
9.93
The causal world model consistently outperforms the baseline in both environments:
The causal model achieves significantly higher success rates, particularly in longer planning horizons. In iTHOR, 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).
Success Rates:
Efficiency: For successful trajectories, the causal model takes fewer steps 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 GrrdWorld, 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 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 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 our experiments can be attributed to the CRI world model's higher accuracy in predicting future states. Both our method and the baseline 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-awan reasoning and planning 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 complex scenarios as CRL and search methods advance. Future work could explore applications to werl-world environments, improve the interpretability of learned causal world models and develop techniques independent of labeled causal variables.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
For reproducibility, we publish the code and models to integrate Causal Representation Learning (CRL) with Language Models (LLMs), as well as the scripts generate data sets used in experiments, on our code repository: implemented using PyTorch (Paszke al. 2019) and PyTorch Lightning (Falcon & The PyTorch Lightning team, 2019). Detailed hyperparameters and dataset descriptions are provided in Section
All models were
and Section 6 with further details in Appendices C
F D
M
and N
In terms of computational resources, all experiments were performed on NVIDIA A100 GPUs. Training the autoencoders takes approximately I to days. Jointly training the normalizing flows and the language heads takes around 0.5 to hour.
ACKNOWLEDGMENTS
The authors thank the Netherlands Organization for Scientific Research (NWO) for their support (VICI grant V1.C.212.053).
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-Hue 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 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 180iv188660010 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, 13:1877-1001,002
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.
11

---

Published as a conference paper at ICLR 2025
Laurent Dinh, Jascha preprint arXiv:1605.0880 2016.
and Samy Bengio. Density estimation using real nyp. arXiv
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.21788, 2024.
William Falcon and The PyTorch Lightning team. PyTorch Lightning, March 2019. URL https
//github.comm/// htt.:///ttttiiintiiaaaaiaeeee A/ghignnng
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. |RL https::/t///tttt.. /ttttp:://openrevieeee net. forum? d=xxxx3xxxxxx3
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 construct and utilize world models for model-based task planning. Advances in Neural Information Processing Systems, (16779081-79099 2023.
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.066 06665, 2024.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint 1803. 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.05997, 2023.
Aapo Hyvärinen and Petteri Pajunen. Nonlinear Independent Component Analysis: Existence and Uniqueness Results. Neural Networks, 12(3):429-433) 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 t plan, but can help planning in LLM-modulo frameworks. In Forty-first International Conference on Machine Learning, 2024.
URL /ettps://openreviie http://wwwww.eeeettnmemeaaaaaaaa
n Eorum? -in--nedd--- EmH
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. arXiv:1312.6114 2013.
Auto-encoding variational bayes.
arXiv preprint
Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International conference machine learning, pp.
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.
12

---

Published as a conference paper at ICLR 2025
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-55
Eric Kolve, Roozbeh Mottaghi, Winson Han, Eli VanderBilt, Luca Weihs, Alvaro Herrasti, Daniel Gordon, Yuke Zhu, Abhinav Gupta, and Ali Farhadi. A12-THOR: An Interactive 3D Environment for Visual AI. arXiv, 2017.
Aneesh Komanduri, Yongkai Wu, Wen Huang. Feng Chen, and Xintao Wu. Sem-vae: Learning identifiable causal representations via structural knowledge. In 2022 IEEE International Conference on Big Data (Big Data), pp. 1014-1023, IEEE, 2022.
Séhastien 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 Prrousi 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:240I. 04890, 2024.
Minne Li, Mengyue Yang, Furui Liu, Xu Chen, Zhitang Chen, and Jun Wang. Causal world models by unsupervised deconfounding of physical dynamics. arXiv preprint urXiv: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. Lim+p: Empowering large language models with optimal planning proficiency. arXiv preprint arXiv:2304. 1477, 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 coordcony 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- arXiv:2208. 14153, 2022.
Weight-Variant Variant Latent Causal Models. arXiv preprint
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 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 Ilm 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: 2019b.
13

---

Published as a conference paper at ICLR 2025
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 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. Hicrarchical Causal Representation Learning. In Causal Representation Learning Workshop at NeurlPS 2023, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Nils Reimers and Iryna Gurevych. Sentence-bert, Sentence embeddings using siamese bert-networks. arXiv preprint arXiv: /0084, 2019.
Jonathan Richens and Tom Everitt. Robust agents learn causal world models. In The Twelfth International Conference Learning Representations, 2024. htttps://oppereeeeeee
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-633 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 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 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 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. / 2022.
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficien model--asee rl through unsupervised object discovery and curiosity-driven exploration, arXiv preprint 1901v 1905.0925, 2019.
14

---

Published as a conference paper at ICLR 2025
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc Le, Denny Zhou, et al. thouuhof-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24877 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 variational autoencoder. arXiv preprint (xxxv:2004.08697 2020.
Dingling Yao, Danru Xu, Sebastien Lachapelle, Sara Magliacane, Perouz Taslakian, Georg Martius, Julius von Kügelgen, and Francesco Locatello. Multi- 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 language models. 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.
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čević, Moritz Willig, Devendra Singh Dhami, and Kristian Kersting, Causal parrots: Large language models may talk causality but are not causal. arXiv preprint arXiv:2308. 1020008..3000 13067, 2023.
Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training. In Proceedings of the IEEEJCVF 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

---

Published as a conference paper at ICLR 2025
A GRIDWORLD ENVIRONMENT
The gridworld environment is dynamic environment of size H x H, where H E 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 E V. obstacles implemented as a color, which differentiates it from other objects within the same entity class.
O, and traffic lights tl E TL. Each entity has fixed corresponding attribute,
The traffic lights positioned in the grid, and each vehicle is facing a specific traffic light. The positions the traffic lights are fixed and immutable, with coordinates (mz,, yu). where Tut Yu E {0, 10.1.... H 1}. Each traffic light has state sn E {red, green). The obstacles have positions (To, No) in the grid, where To E {0, 1,...... H - 11. and these positions can only change through interventions performed
them. The vehicles have positions (Ir: Me) in the grid, where Tep y/ {0, 1.... H l}. and an orientation 02 {up, down, left, right}. The vehicle positions change according to the following dynamics:
Let e be vehicle at position (I Uo) with orientation 0,, associated with a traffic light tl at position (ru, Mu). We say that the vehicle facing the traffic light tl if and only if one of the following conditions is satisfied:
1. 0, up and v. Tat and y/ yu
2. 0, down and H, = zu and yo < ys
3. 0g = left and yo = yo and we > Tt
4. 0g right and ye = Yul and Te Fel
If the vehicle is facing the traffic light tl, it will move forward the cell (x% v() the next timestep if and only if all of the following conditions are satisfied:
1. The traffic light il has state of green, i.e., 81 green.
2. There obstacles in the cell (ar v.),, i.e., o E O (Zo: (1h v())
3. There are traffic lights in the cell (o v(), i.c., tl TL (In yu) = (I' v/)).
4. The cell (xt,, v()) is within the grid boundaries, i.e., 0 S x < and 0 S v. H.
The new position (z', v.) is determined by the vehicle's current position (æe. Uo) and orientation 0g as follows:
(Eer yo - 1)
if A, = up
(Fv, yo + 1)
if 0, down
(x vo))
(x. Yu)
if A, left
(xr + 1,y0)
if 0, = right
first, a step in the environment dynamics is executed: then, 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 single-cell displacements; for obstacles, the direction is stochastic, while for vehicles, it is deterministically forward. Vehicle intervention further constrained by the absence of obstacles or traffic lights in the target cell, adherence to environment boundaries, and the corresponding traffic light displaying red signal. no-operation (NOOP) intervention also permissible. This tripartite sequence sequence-envinee environmental progression, intervention, and state documentation constitutes complete intervention eycle. These interventions correspond to regime variables R', which are then represented using natural language.
Interventions The intervention process follows specific sequence:
Causal Variables The causal variables the gridworld environment are the positions of the vehicles (Fa, (w), the positions of the obstacles (Io; yo), and the states of the traffic lights su-
16

---

C(r.). The generation process for Published as a conference paper at ICLR 2025
B ITHOR KITCHEN ENVIRONMENT EMBBDDED AI
The iTHOR (Kolve et al., 2017) kitchen environment is based on the FloorPlan dataset, featuring 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 interactive objects f E F, and receptacles E R. Movable objects include a plate with potato and an egg. Fixed interactive objects comprise microwave, stoves, cabinet, and toaster. Receptacles include the counter, microwave (when open), and pan (for the egg). Each object has state So S.. where So is the set possible states for object o. For binary state objects (e.g,, microwave, cabinet), S, = open, closed or active, inactive. For movable objects, So includes their position (Tim, Vm: Fm) the space and binary pickup state. The set of possible actions includes:
E M, fixed
ToggleObject(o): For o E (microwave, stoves, toaster]
OpenObject(o): For o E (microwave, cabinet)
PickupObject(mn) For m M PutObject(m,, r): For E M,r E R
MoveObject(m): For m E
NoOp: No action performed
The availability of actions depends on the current state of objects. For example:
ToggleODjcc((microwwee) is valid iff micnowave = closed
(4)
OpenOhject(mierowave) is valid Smicrowave inactive
(5)
The regime variable R E [0, 1 represents the normalized click-location on the image to select the object for interaction. Let Io the set of pixels belonging to object o in the current frame. Then:
1
R
H W (Ir,)), where (xr, v) Uniform(Io)
where H and W are the height and width of the frame respectively. The causal variables C C....., Cc in this environment correspond to the states and positions of objects. Binary variables (e.g., Cabinet- Open. Microwave-Activer take values in 0, 1, while position variables (e.g.. Egg--os-x/ take continuous values in [0, 1]. normalized to the environment's dimensions. Observations are generated as high-resolution images X E 1512x512x3 then downsampled to X E I2 25662566x using bilinear interpolation.
C TEXT-BASED ACTION REPRESENTATION GENERATION
C.1 GRIDWORLD ENVIRONMENT
For the GridWorld environment, we implement probabilistic context-free grammar (PCFG). The
PCFG includes:
A set of adjectives A, for each object type o E O = (traffiic light, vehicle, obstacle)
A set of action modifiers M
A set of action verbs V, for each action type A {mooe,, turn, change state]
Let C 3 + Eg be function mapping RGB values to finite set of color names Σ., For each object with RGB value To, we compute its color name an action a on object o can be formalized as:
D(a, o) = m V. the adjo Co
(7)
where m ~ P(M), Va E Va. adjo ~ P(A.), and P(.) denotes the probability distribution defined by the PCFG. Example: Consider action move blue car to the right. Let r. (0. 255). C(r.) = "blue", m "skillfully". D2 = "moved", and adj. = "sleck". The generated description would be:
(move right, car) "You skillfully moved the sleek, blue the right."
17

---

Published as a conference paper at ICLR 2025
C.2 ITHOR ENVIRONMENT
For the iTHOR environment, we define a mapping function f A x O + Σ, where is the set of possible actions, O is the of objects, and Σ is the set of all possible strings over the alphabet. Let T. A + V be a function that maps actions to verb phrases, and T, O + E" be function that maps objects to descriptive phrases. The generation process for an action a on object can be expressed as:
(9) Example: For the action of opening microwave, let OpenObject and Microwave. As- sume To (OpenObject) = "adjusted" and (Microwae = "the microwave's s door". The generated description would be:
f(a. o) You (a(a) To(
f(OpenObject, Microwave) You adjusted the microwave's door."
(10)
C.3 TOKENIZATION AND INTEGRATION
Let Σ- N be a tokenization function that maps string a sequence of 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 resulting in the final representation t' N'. For a trajectory of actions ai.... On on objects 01, Ons we generate sequence of tokenized descriptions t,, ln where:
t = pad( (D(a, 0.)), for GridWorld
(12)
4/2 pad( o)),, for iTHOR
(13)
D DATA GENERATION AND PREPARATION
For each environment, we generated multiple datasets as shown in Table
Table 5: Dataset specifications for each environment
Dataset
Size
Description
Training
10000 trajectories of 100 steps
Used for model training
Validation
1000 episodes 100 steps
Used for model validation
Test
1000 episodes of 100 steps 100 episodes of 100 steps
Used for final evaluation
ICL
Used for in-context learning
N-step evaluation
100 episodes of 100 steps each,
Used for N N-step experiments
for each N value
D.1 DATA GENERATION PROCESS
The data generation process for both environments follows these steps:
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 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 for the desired number of steps (100 in case).
18

---

Published as a conference paper at ICLR 2025
For the Gridworld environment, valid actions include toggling traffic lights and performing 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 putting down movable objects, and performing no action.
For N- N-step experiments, generate multiple datasets, each corresponding to different value of N:
Gridworld: We create separate datasets for N E (2,4, 6, 8).
. iTHOR: We create separate datasets for N E (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 a 3-tuple (ini- tial pilssttteccusalll variaabes, actions, end_ und_state___ saatt__cusal___ variables). For each iteration during training or evaluation, we randomly sample two examples from this pool to provide context for the model. This process similar 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 examples covering every aate-action pair at least once. For 4-step experiments, we use examples covering at least 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 I-step trajectories. We select the actual action taken in the environment for "good" evaluations, providing positive example of correct action. For "bad" evaluations, 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 jaatee following the process described Section
E. AUTOENCODER
The autoencoder trained from scratch using 10 times more samples than the main dataset to ensure robust representation. This approach is justified by the relative ease of obtaining unlabeled, random samples from an environment. pretrained image representation model can be employed by adding a learnable linear projection the required dimensions and training with the original dataset size.
scenarios where this not feasible, transfer learning from
For the Gridworld environment, we implement an autoencoder with 40 latent dimensions and 64 hidden channels. Both the encoder and decoder consist residual blocks with SiLU activation functions. We incorporate the CoordConv operator (Liu al.) (018) to better capture coordinate information from images. For the iTHOR environment, we employ the autoencoder architecture from BISCUIT (Lippe al. 2023)
19

---

Published as a conference paper at ICLR 2025
E.2 NORMALIZING FLOW AND TRANSITION MODEL
For both the normalizing flow and transition model, we use the same architectures and hyperparam- in BISCUIT (Lippe et al. 2023) 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 a pretrained Sentence Transformer (Reimers & Gurevych, 2019), specifically the all-MiniL.M-L6-v22 mode augmented with a 2-layer MLP head with 64 hidden dimensions. For iTHOR, we use pretrained SigLIP model (Zhai et al. 2023 with a similar 2-layer MLP head. In both cases, the pretrained encoders remain frozen during traiiing with only the MLP head being updated.
E.4 TRAINING PARAMETERS
Key training parameters for each environment are follows: For Gridworld, we use learning rate of 10 for the main model and 10-3 for the MLP, batch size of 384, and train for 300 epochs. For iTHOR, we use learning rate of x 10-3 for the main model and 3 x 10-3 for the text MLP, batch size of 64, and train for 100 epochs. Both environments employ warmup period of 100 steps and 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 tross-validation to select the optimal hyperparameters for MIP 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 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 exploration weight parameter in the search algorithm due to the interaction between exploration- exploitation trade-off and reward scaling.
G CAUSAL MAPPER
The causal mapper my extracts interpretable causal variables from the learned disentangled represen- tations.
This process allows for non-injective mapping from latent dimensions to causal variables. For instance, we have a causal variable "cabinet_statee statee, the first stage might learn that latents 1, 5, and 7 are the most predictive for this variable. In the second stage, specific predictor would leam to map from these dimensions to either 0 (closed) or 1 (open).
The causal mapper my is implemented in two stages:
G.1 TARGET ASSIGNMENT
This stage uses a single MLP fussign to predict all causal variables from each latent dimension independently:
1--1---22
https::/hugginggaaas ht
20

---

Published as a conference paper at ICLR 2025
Č Jussign (a M)
(14)
where M {0, is mask and elem multiplication. M to the identity.
For each latent dimension i, we create a mask M; where only the i-th element is and the rest are O. We then batch these masks along with the corresponding masked latent vectors:
M1 O
M1 M2
M
(15)
2 M
LM
This batched input is fed into Jussign- which outputs predictions for all causal variables for each masked input. The output shape is [L. CJ. where L is the number of latent dimensions and the number of causal variables.
We then compute the correlation between these predictions and the ground truth causal variables. This allows 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 relevant for each causal variable to determine each My.
G.2 CAUSAL PREDICTION
Individual MLPs feauall, are trained for each causal variable j. using only the relevant latent dimen- sions identified in stage 1:
Õ = feausal S M2)
(16)
where Mj is the mask for causal variable j.
The output layer of each feausal.j is adjusted based on the priori known type of the causal variable (categorical, numerical, angle).
H STATE DESCRIPTION GENERATOR
The state description generator 8 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 operate either stochastically or deterministically, depending the application's needs:
1. Stochastic:
This approach a language model with temperature greater than which allows for a variety of possible descriptions for the same state. This variability can be useful in scenarios where diverse language outputs desired.
2. Deterministic: This method involves either setting the temperature of a language model to O, ensuring consistent outputs, or using rule-based system that directly maps causal variables to fixed phrases or sentences.
H.2 EXAMPLE OF STATE DESCRIPTION GENERATION
For instance, given dictionary of causal variables follows:
21

---

Published as a conference paper at ICLR 2025
"cabinet_ae :esiinet_state:: 1, "light [eiggt_collo") 0, "door_angle": 45
The state description generator might produce sentence like: "The cabinet is open. The traffic light showing a red signal. The door is partially open at 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 reason effectively about state transitions.
The modular nature of our framework allows for easy substitution 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, implement conversion of visual states to natural language descriptions using the ground truth causal variables. This process ensures fair comparison with causal world model while maintaining the LLM's ability reason about the environment.
For each initial state, we extract the ground truth causal variables and the same rule-based state description generator employed in our causal world model to convert these variables into natural language. For example, in GridWorld, state with causal variables blue aae_ear_____ car_aaa x 2, blue_cc _ae__cr____ y: 3, cyan light _tatee green would be converted to "The blue car is position (2.3). The cyan traffic light showing green.".
The baseline LLM then uses this initial state description to reason about subsequent states and actions, relying on its world model capabilities 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 absolute error (MAE) across all dimensions against training set size. Standard deviation bands are computed from three independent training illustrate the variance in performance across different training instances.
22

---

Published as a conference paper at ICLR 2025
2. Dimension-wiss 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.
Ovcrall Perbermsse
Cousal Mapper Dimension Evolution
40.0
6:18
6:4
ins.iciitii S 255 2551 HEI
0.10
tudking 000, 300, MAA
000
105
Stabung
(a) Model performance (MAE) vs training size the causal mapper in GridWorld. Error bands represent one standard deviation across three independent runs.
(b) Dimension- wise performance analysis showing 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-wee 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 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 predicting traffic light states even 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 both the Gridworld and iTTHR 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, make the following adjustment:
23

---

Published as a conference paper at ICLR 2025
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 learn and evaluate other than the fact that it moved not.
The boulder's position can vary due to random factors not captured in our causal model. By excluding it from our evaluation, focus the aspects of the environment that causally determined by the actions and states we' 'e modeling.
Rationale:
K.2 ITHOR ENVIRONMENT
For the iTHOR environment, implement more nuanced approach:
Coordinate Categorization: We categorize the x, y, and = coordinates for objects with stochastic movements into discrete position categories.
Category-basee Evaluation: Instead of comparing exact coordinates, we check whether objects end up the correct category positions after an action.
Rationale: In iTHOR, object movements can have slight variations due 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 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. 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 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 World samples, with 10 runs per sample after warmup, using single NVIDIA A100-40GB GPU.
Our CRL world model consists of three main components: an autoencoder (4.5M parameters), normalizing flow (2.9M parameters), and 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 reprennrpiessntiing an approximately 82x speedup. This computational difference has significant implications for planning tasks. For example, 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 necessary. The efficiency of our CRL world model enables more extensive tree searches and faster iteration during planning, while maintaining high prediction accuracy demonstrated in main experimental results.
24

---

Published 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 sequence of latent states and their corresponding natural language descriptions.
Algorithm Inference with the Causal World Model
Require: Observation space X, latent space 2, action description space L. action encoding space A; observation encoder cvs normalizing flow fa. action encoder Le, transition model Pa, causal mapper mg. state description generator s; initial observation Xo A;; action descriptions
LT
1: function E(X E 2)
2: E Lo(ew (X))
b Causal encoding of observation
3: return E
4: end function
5: function ENCODEACTION(L L)
6: a Lc(
> Encode action description into action latent space
7: return
8: end function
9: function G(z E 2)
10: C+ (2)
> Map latent state to causal variables
++ sCC)
> Generate natural language state description
12: return L
end function
14: function SAMPLENENTSSATTESS E 2, 10EEA A)
15:
Zt+1 ~ Pw (2+41 ze, a()
Predict next latent state
16:
g(2)
Generate current state description b Generate next state description
17:
4+11 + g(z++1)
18:
return (*++1, lu l++1)
19:
end function
20: function INFERENCETRAFRRTRRYCCC (T117)
21:
2o - {(Xo)
Initialize latent state Generate initial state description
Eo - ((20)
23:
yield (zo, Lu)
for t 0o - do
25:
at f- EEccodeAction('")
Encode action description
26:
(z++1, (x+1)) +- SAMPLENEXTSTATT(2...C yield (2441,141)
27:
28:
21 4- 21+1
Update latent state for next iteration
29:
end for
30:
end function
25

---

Published as a conference paper at ICLR 2025
N MODIFIED MCTS PLANNING ALGORITHM
We adapt the Reasoning via Planning (Ha et al. 2023) Monte Carlo Tree Search (MCTS)
algorithm (Koesis & 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 X causal world model (Algorithm 1. LLM agent, depth limit L, number of roll-outs N, exploration weight intuition ICL samples Diat self-evaluation ICL samples
Diel
Initialize memory of actions A ~ L. children e: 2 + 2 and rewards r: x L TR
2: Initialize the state-action value function Q 2 L m IR and visit counter N 2 + N
3: z" < E(X),, e (z2")
Initialize root node
4: for + ..... - do
5: t 0, + z + g"
6: while N(22) > do
Selection
7:
N(22') + N(z')) +
8:
0r -- arg max max
Q(zz, + WV ((c(x'_aa)
Mix')
9:
t - r(z', az). gl+1 + clz', ar)
+
zt +11, G(z')
11:
end while
while x' not a terminal state A S do
13:
A <- GetValidActions(("
for E A do
> Expansion Use the Causal World Model
15:
z-1 - gt+1 +- SAMPLENEXTSTATEE2 La)a
16:
F Vintuition + log PLLM(a| EL, Diet
F evllFeval +- log (ELLMM"M"ooo aa Didl)
18:
r(z', a) f- Fintuition Tsetf-oval
19:
Update A(z') + A(z') (a}) clzzzzooooo + ++1 end for
20:
21:
04-1 f arg max r(z))r(z';,
b Simulation
22:
Ii f- r(z', (in',ap 1,+1), 1--1 +- c(z', a(41) z111, + G(z')
+
24:
end while
25:
for tt..,, do
Back propagation
Update Q(z' ,a) with fre re41....,
27:
end for
28:
end for
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 X (line 3).
2. Causal World Model Integration: We employ trained causal world model (Algorithm 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 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 sample actions directly from the LLM for open-ended domains where the action space is not easily enumerable.
26


Reviews:
[{"review_id": "oLNEw4Ntkb", "reviewer": "Reviewer_Uz31", "summary": "The paper introduces a framework that combines LLMs with CRL to improve reasoning and planning tasks. The framework leverages both LLMs' common sense knowledge and CRL's \"look-ahead\" ability on causal structures in environments. During inference, LLMs use a causal world model as a simulator for generating and processing action and state descriptions in natural language. This can be well integrated with MCTS algorithm from previous work. Tests on causal inference and planning tasks reveal that this causally-aware approach outperforms traditional LLM methods in complex, long-term planning scenarios.", "questions": "Questions:\n- How do you make sure the objects in the image X are identified with the Auto Encoder? if some objects are missing the casual model cannot run properly right?\n- How does the baseline RAP work with LLama? LLama 8B does not take vision input.\n- L285: the description is confusing, why causal mapper is trained using image input?\n- Is LLM baseline just RAP? should just say RAP it in the table. Baseline LM is super confusing, there are many other simpler LLM prompting baselines in planning.\n\nTypos:\n- L269: p_\\phi(E_t | z_t) not p_\\phi(X_t | z_t)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The integration of CRL and LLM planning is novel and interesting, it is straight forward to integrate it with multiple other LLM-based search algorithms, not only RAP.\n- The paper investigates the form of action representation in the casual world model, and provides detailed results on them.", "weaknesses": "-  Reasoning via Planning is not the sota method, [1] combines LLM as a world model and LLM as a policy model, with MCTS, should be better than RAP\n- Though not discussed in RAP, other literature [2] [3] suggest increasing the number of LLM calls will substantially improve the search algorithm results. How is the overhead and accuracy of LLM+CRL planning vs. increasing LLM calls with tree-search or iterative prompting?\n- I'm concerned about the ability of the system to work on more complex visual environment. For example, if the accuracy of the casual world model on look-ahead is limited, we should still rely on executing the actions in the real simulator and extract observation there. Related to this concern, a comparison with RAP+the real simulator execution (some oracle) is missing.\n\n[1] Zhao et al., Large Language Models as Commonsense Knowledge for Large-Scale Task Planning, 2023\n\n[2] Yao et al., Tree of Thoughts: Deliberate Problem Solving with Large Language Models, 2023\n\n[3] Zhang et al., ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search, 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2V3pWH3vPo", "reviewer": "Reviewer_wNeN", "summary": "The paper proposed to learn a causal world model from a sequence of states and actions, where the states are images and actions are described in natural language. The proposed method builds upon BISCUIT and proposes representing the actions in natural language. The paper conducted experiments in GridWorld and iTHOR environment and the R^2 scores of the learned representation are claimed to be higher in low-data regimes. Experiments also show that the learned causal world model is more accurate at predicting the next state, and improves planning performance when compared to a general LLM (Llama 3).", "questions": "1. How is coordinate-based action representation implemented? A 2-d vector, or simple text like “(2,3)” encoded by the encoder?\n2. What does HB action representation look like and how does it differ from TB? An example of CB, TB, and HB would help to explain.\n3. line 377, “better sample efficiency” would allow TB to perform well in extremely low-data scenarios as well, the reviewer think the experimental results doesn’t support the claim of “better sample efficiency” of TB.\n4. In table 2, the baseline language model and the causal world model both predicts the next state in natural language, how is it evaluated against the ground-truth next state?\n5. How big is the set of annotated images used to train the causal mapper m_\\theta?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper proposed a method to learn a causal world model from a sequence of image states and text action descriptions, and demonstrated superior performance in the accuracy of the learned world model.\n- Presentation of the paper is clear and easy to follow.", "weaknesses": "1. Some components of the framework would not be available in more realistic environments, e.g. 1) a set of annotated images with ground-truth causal variables is used in training, which is likely not available as we may not know the causal variables for more realistic environments; 2) a rule-based state description generator may not be available for complex environments where we don't know what are the true causal variables.\n2. Given the simplicity of the environments, and that the proposed method is trained on these particular domains while the baseline is a general LM, the superior performance of the learned causal world model is less convincing. I would suggest comparing with a supervised fine-tuned version of the baseline LM. Since the proposed method uses a set of annotated images to train the causal mapper, supervised fine-tuning is possible with these annotated data: use the state description generator to convert the ground-truth causal variables of the states to natural language, and we can obtain a sequence of natural language action and natural language description of the states.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PFTAMTVpx0", "reviewer": "Reviewer_jsxQ", "summary": "The paper proposes a novel framework that integrates Causal Representation Learning (CRL) with Large Language Models (LLMs) to enhance reasoning and planning capabilities in interactive environments. The framework utilizes a causal world model that maps high-dimensional state representations to causal variables, which are linked to natural language expressions. This integration allows the LLM to interact with the causal world model, effectively simulating multiple possible futures before taking actions. The approach is evaluated on causal inference and planning tasks across different environments and temporal scales, demonstrating that the causally-aware method outperforms traditional LLM-based reasoners, especially for longer planning horizons.", "questions": "Can the authors elaborate on how the framework handles situations where the causal factors are not easily identifiable or where the causal relationships are highly complex?\n\nFollowing on the above point, how were the parameters chosen for the various components of the framework, and how sensitive are the results to these parameters?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The exploration of text-based action representations is particularly interesting, as it shows potential advantages in low-data regimes. I particularly liked the connection to the RL decision making problem. The results indicate that the proposed framework can potentially improve performance in causal inference and planning tasks, which is valuable for the broader ICLR community.", "weaknesses": "Please see the questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Mfivtt6cf1", "reviewer": "Reviewer_ZjUs", "summary": "The author proposed to introduce causal representation learning (CRL) to better understand the causal structure of a environment, and they incorporate such CRL with LLM to enable causally-aware planning by mapping causal representations to natural language.", "questions": "+ Do we have examples of decoded text (from CRL to LLM)? Is that human-readable? Is there any possible insights with analysis of the decoded text?\n+ (Clarification) In my understanding, the performance increase comes from (1) the text decoded from CRL makes it easier for LLM to reason and plan (2) the CRL model has a higher accuracy in predicting future states. If my understanding is correct, do we have ablation experiments on the effect of the two aspects?\n+ (Clarification) Does the Baseline LM refer to RAP + LLaMA 3 or LLaMA3 with other prompting methods?\n+ In Table 1, **TB** outperforms **HB** starting from 1.0%, outperforming **HB** much at **1.5%**, but performs weaker than **HB** in 100%. What are the trend between 1.5% and 100% subsample rate? Does it mean the necessity of **TB** decreases when the data scales up enough?\n+ The **HB** settings shall contain all information in the **TB** settings. Why would the model performs worse than **TB** with increased subsample percentage? Since **HB** performs better in 100% ($10^6$ image states), would it be better to use **HB** in subsequent experiments?\n+ Can you include SNA (success weighted by number of actions) in the planning results? It's a common metric to combine success and # of steps.\n+ (Typo) Page 4 footnote, prioi -> prior\n+ (Typo) Figure 1, missing an arrow from $\\bf C^1$ to Text Decoder.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposes to introduce CRL for world model to better understand the causal structure of environments. This intuitive idea provides a solid future approach to further improve the performance of world models, which is essential to build with intelligent agent with System 2 reasoning abilities.\nThe author analyzes the underlying math principles of the proposed CRL model and shows with experiments of the performance in metrics of comparison between induced state variables and ground-truth, and multi-step high-level accuracy of the predicted states. The experiments show large performance increase than baseline pure language model.\nThe author also use the result causal world model to do experiments on synthetic planning tasks. The results show consistent superior performance to baseline LM, especially on longer planning horizons.", "weaknesses": "+ It's not clear how much the text decoded from CRL model helps the LLM reasoning. It would be clearer if the author can show some examples and/or conduct ablation experiment on this. (also see questions)\n+ The experiments compare **CB**, **TB**, and **HB** are based on the low-data scenarios, but the decision to use **TB** in subsequent experiments is for 100% data. This decision lacks a continuity of experiment settings, and it's unclear which setting would be the best for subsequent experiments. (also see questions)\n+ The paper is a good starting of causal world model studies, but I have to mention that the experiment environment is synthetic and relatively easy. It would be more fascinating to use more general environments, and there leaves doubt whether the method can be easily scaling up to more general or even real world scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)