Language Agents Meet Causality -- Bridging LLMs and Causal World Models

Thank you to Reviewer ZjUs for the comprehensive feedback. Your detailed questions about text representation and experimental design have helped us clarify the presentation. --- ### **Weakness:** > "It's not clear how much the text decoded from CRL model helps the LLM reasoning." The decoded text is human-readable and follows a consistent format. Examples include "The blue car is at position (2,3)" and "The cyan traffic light is green." The text format enables the LLM to leverage its pre-trained knowledge while maintaining the interpretability of the state representations. This allows for easier debugging and verification of the model's reasoning process (also see the answer to the corresponding question). --- ### **Weakness:** > "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." ### **Question:** >"Why would [HB] perform 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?" Section 6.1 is intended to serve two purposes: (1) demonstrating that TB representations are effective for CRL, and (2) showing they outperform coordinate-based representations in low-data regimes. We first establish that text-based representations are effective and then use the strongest CRL module from those. Thus, for planning, text-based representations are preferred as coordinate outputs would be dissimilar to the LLM's pretraining distribution. Given enough data, the evaluated representations converge to practically indistinguishable performance. Additionally, our text-based representations offer important advantages: - **Robustness to paraphrasing** – different descriptions of the same action map to the same latent space region. - **No need to enumerate all possible actions across environments.** - **Independence from specific action encoding formats** used during pre-training (e.g., third-person vs first-person descriptions, varying adjective usage). These properties arise from our language encoder backbone's ability to map varied textual descriptions to consistent latent representations, making our approach more practical for real-world deployment. While HB performs well with full data for pure representation learning, TB is more suitable for the complete framework due to its natural alignment with LLM capabilities. --- ### **Weakness:** > "The paper is a good starting of causal world model studies [...] but the experiment environment is synthetic and relatively easy. It would be more fascinating to use more general environments [...] or even real world scenarios." Our framework is designed to benefit from advances in CRL, which is a rapidly developing field. As stronger CRL methods are being developed for handling more complex environments, they can be directly integrated through our modular design. --- ### **Question:** > "Do we have examples of decoded text?" Yes, the decoded text is human-readable and interpretable. For example, it might output: "The blue car is at position (2,3), while the traffic light is green, …". This provides insights for failure analysis by allowing us to discern whether errors stem from incorrect state representations. --- ### **Question:** > "(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?" The performance gain stems from (2). Both approaches use identical state representations, achieved by maintaining consistent text formatting in the in-context learning examples for our method and the baseline. The LLM in the baselines then adheres to the same format when generating state representations. This makes it possible to attribute the performance improvements to the CRL world model's superior state prediction accuracy. We clarified this in lines 521–526. --- ### **Question:** > "(Clarification) Does the Baseline LM refer to RAP + LLaMA 3 or LLaMA 3 with other prompting methods?" It refers to RAP+LLaMA 3, now clarified on line 353 and in tables 2, 3, and 4. Direct prompting for plan generation performed significantly worse, even with additional context examples. Our framework is orthogonal to the specific search algorithm used. --- cont.

We value Reviewer jsxQ's questions about framework limitations, which helped us better articulate our parameter selection process and robustness considerations. --- ### **Question:** > "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?" When the CRL component correctly identifies the causal factors, our framework enables effective reasoning and planning, as demonstrated by our results. In cases where causal identification is challenging, the framework's performance would be limited by the upstream CRL component's capabilities. --- ### **Question:** > "How were the parameters chosen?" The hyperparameters were selected as follows, now detailed in Appendix F: - **Causal encoder:** We use the hyperparameters from the BISCUIT paper. - **Text encoder:** Cross-validation across 5 folds. - **Search parameters:** Bayesian optimization with 15 trials. Our experiments throughout the project indicated low sensitivity to model training hyperparameters but higher sensitivity to the search exploration parameter due to reward scaling.

We thank Reviewer wNeN for their thorough analysis. Your inquiries about action representations and experimental design have helped us bolster our presentation. --- ### **Weakness:** > "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." Our framework requires annotation only of causal variables and their values (e.g., object positions, light states), not their relationships or temporal dependencies. This is significantly easier than annotating temporally consistent sequences of states with their corresponding actions. The causal mapper requires a small labeled set (6.25% of validation data: 1000 frames for GridWorld, 1500 for iTHOR), used only for the decoder phase. We added an analysis of the mapper’s data efficiency to the paper (see Appendix J). While we used a rule-based state description generator for simplicity, this component is modular and can be replaced with a sequence-to-sequence model or fine-tuned LLM for more complex environments. We have extended Appendix H.3 to detail these alternatives. The key requirement is consistency in the mapping from causal variables to descriptions, achievable through various approaches. --- ### **Weakness:** > "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." Thank you for the suggestion! Fine-tuning the baseline LM would require triples of state, action, and next state in *natural language* to learn the causal dynamics. This is a different experimental setup that makes a stronger assumption on the available annotations than our method. The labeled examples we use for the causal mapper are unordered individual frames (clarified on lines 284-286), lacking the temporal structure needed to learn causal dynamics. Our approach learns the causal dynamics from the environment with CRL. In addition, the benefits of supervised fine-tuning could be limited because: 1. **2-9 planning problem solutions spanning multiple steps are already used for in-context learning depending on the experiment**, detailed in Appendix D.2. 2. **The pre-prompt and few-shot examples already provide environment-specific information.** --- ### **Question:** > "How is coordinate-based action representation implemented? A 2-d vector, or simple text like '(2,3)' encoded by the encoder?" It uses fixed 2D sinusoidal encodings mapping 2D inputs to a higher-dimensional space using high-frequency functions, similar to the non-learnable components in NeRF/Fourier feature encodings [1, 2]. This is now clarified in lines 328-333. --- ### **Question:** > "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." We have added examples on lines 328-333: - **CB:** `(2,3)` → sinusoidal encoding - **TB:** `"move two steps right and three steps up"` → text embedding - **HB:** concatenation of both above representations --- ### **Question:** > "line 377, 'better sample efficiency'..." Thank you for this observation. We have weakened the claim about sample efficiency on line 377. --- ### **Question:** > "In table 2, how is it evaluated against the ground-truth next state?" We compare the predicted states against ground truth by focusing on causally relevant aspects while accounting for environmental stochasticity. For instance, in iTHOR, we categorize object positions into discrete regions (e.g., “on counter,” “in microwave”) rather than comparing exact coordinates, since physics-based interactions can lead to slight positional variations even for identical actions. This ensures we evaluate meaningful causal predictions rather than exact numerical matches. Full evaluation details are provided in Appendix K. --- ### References [1] Tancik, M. et al. (2020) *‘Fourier features let networks learn high frequency functions in low dimensional domains’,* Advances in Neural Information Processing Systems, 33, pp. 7537–7547. [2] Mildenhall, B. et al. (2020) *‘NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis’,* in ECCV.

We appreciate Reviewer Uz31's detailed feedback, which helped us clarify our framework's relationship to existing planning methods and computational advantages. ### **Weakness:** > "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" We chose RAP as a representative baseline to isolate the impact of causal understanding, as our framework is agnostic to the specific planning algorithm used. We have clarified the statements on RAP’s classification as SoTA and our framework’s flexibility with respect to the search algorithm (lines 353-356). --- ### **Weakness:** > "[...] [2] [3] suggest increasing the number of LLM calls will substantially improve [RAP’s] results. How is the overhead and accuracy of LLM+CRL planning vs. increasing LLM calls with tree-search or iterative prompting?" We performed benchmarks comparing single-step predictions between the LLM-based world model (LLaMA 3 8B via ExLlamaV2 quantized to 6 bits) and the CRL world model on 5 GridWorld samples, with 10 runs per sample after warmup on a single NVIDIA A100-40GB GPU. The CRL module – totaling 36.1M parameters – achieved 27ms/inference compared to the LLM's 2.2s/inference; an ~82x speedup. This computational difference has significant implications for planning: a 10-branch tree search takes ~22s with LLM calls versus ~0.27s with CRL. Results are detailed in Appendix L. --- ### **Weakness:** > "I'm concerned about the ability of the system to work on more complex visual environments [...] we should still rely on executing the actions in the real simulator and extract observation there" Assuming access to an accurate simulator is a strong requirement, feasible only in controlled, simulated environments. In real-world settings, this would equate to executing actions directly in the environment, where some actions are irreversible (e.g., once an egg is cracked, it cannot be reverted to make a hard-boiled egg). Our method learns the causal dynamics to avoid performing irreversible actions in the world. Furthermore, CRL is a rapidly advancing field, and our framework will directly benefit from these advances through its modular design. As stronger CRL methods emerge for handling complex environments, they can be integrated via the causal encoder component. --- ### **Question:** > "How do you make sure the objects in the image X are identified with the Auto Encoder?" Our method shares the same assumptions as BISCUIT regarding the autoencoder's ability to capture objects in its representation. Specifically, we assume the autoencoder can achieve perfect reconstruction, which is a common assumption in CRL methods [4, 5, 6, 7, inter alia]. If the autoencoder fails to capture some objects, this negatively affects downstream performance. Nevertheless, the autoencoder module is orthogonal to our method and can be replaced with stronger alternatives if needed, now clarified on lines 272-275. In practice, while perfect reconstruction might not be achievable, modern architectures can achieve high-quality reconstructions sufficient for learning meaningful representations, as demonstrated by our empirical results. --- ### **Question:** > "How does the baseline RAP work with LLama? LLama 8B does not take vision input." To enable LLaMA to process the environment, we provide it with a language description of the ground truth causals for the initial state (using the rule-based state description generator with the ground truth causal variables of the state as input). This experimental setup is now detailed in Appendix I. --- ### **Question:** > "L285: why causal mapper is trained using image input?" The causal mapper is trained on the output of our causal encoding pipeline (autoencoder and normalizing flow). That is, the training data of the causal mapper are pairs ($\mathbf{z}$, $\mathbf{C}$), where $\mathbf{z}$ is a vector of the identified and separated causal variables from the normalizing flow and $\mathbf{C}$ are the ground truth causal variables. We have clarified this process on lines 284-286. --- ### **Question:** > "Is LLM baseline just RAP? Should just say RAP it in the table." Yes, the baseline is RAP, thanks for the suggestion. Now clearly stated in tables 2, 3, and 4. ---

Thanks to the authors for the reply. The author addressed most of my concerns. Though I still think a perfect visual reconstruction is a too strong assumption (esp for complex env), as it is a common CRL issue, I don't think it's a hard negative of the paper. Overall, I believe the merits of the paper outweigh its flaws. I'll keep my score.