# Better Decisions through the Right Causal World Model

arXiv: https://arxiv.org/abs/2504.07257 · HTML: https://mlsanity.com/paper/2504.07257/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 15 of its 19 references that are in the map.
- First posted on arXiv 9 Apr 2025.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)

## Builds on (most cited first, 15 of 15)

- [Fast Segment Anything](https://mlsanity.com/paper/2306.12156/) (2023)
- [Unsupervised State Representation Learning in Atari](https://mlsanity.com/paper/1906.08226/) (2019)
- [Making deep neural networks right for the right scientific reasons by interacting with their explanations](https://mlsanity.com/paper/2001.05371/) (2020)
- [Goal Misgeneralization in Deep Reinforcement Learning](https://mlsanity.com/paper/2105.14111/) (2021)
- [Right for the Right Concept: Revising Neuro-Symbolic Concepts by Interacting with their Explanations](https://mlsanity.com/paper/2011.12854/) (2020)
- [An Investigation into Pre-Training Object-Centric Representations for Reinforcement Learning](https://mlsanity.com/paper/2302.04419/) (2023)
- [Interpretable and Explainable Logical Policies via Neurally Guided Symbolic Abstraction](https://mlsanity.com/paper/2306.01439/) (2023)
- [OCAtari: Object-Centric Atari 2600 Reinforcement Learning Environments](https://mlsanity.com/paper/2306.08649/) (2023)
- [Interpretable Concept Bottlenecks to Align Reinforcement Learning Agents](https://mlsanity.com/paper/2401.05821/) (2024)
- [Boosting Object Representation Learning via Motion and Object Continuity](https://mlsanity.com/paper/2211.09771/) (2022)
- [End-to-End Neuro-Symbolic Reinforcement Learning with Textual Explanations](https://mlsanity.com/paper/2403.12451/) (2024)
- [HackAtari: Atari Learning Environments for Robust and Continual Reinforcement Learning](https://mlsanity.com/paper/2406.03997/) (2024)
- [Bad Habits: Policy Confounding and Out-of-Trajectory Generalization in RL](https://mlsanity.com/paper/2306.02419/) (2023)
- [SPARTAN: A Sparse Transformer World Model Attending to What Matters](https://mlsanity.com/paper/2411.06890/) (2024)
- [Towards Generalizable Reinforcement Learning via Causality-Guided Self-Adaptive Representations](https://mlsanity.com/paper/2407.20651/) (2024)

## Built on by (most cited first, 3 of 3)

- [GRAIL: Autonomous Concept Grounding for Neuro-Symbolic Reinforcement Learning](https://mlsanity.com/paper/2604.16871/) (2026)
- [Causal Reward World Models: Zero-shot Reward Design for Automated Skill Generation](https://mlsanity.com/paper/2606.23280/) (2026)
- [What is the goal of unsupervised machine learning?](https://mlsanity.com/paper/2610.11697/) (2026)

## Abstract (arXiv)

Reinforcement learning (RL) agents have shown remarkable performances in various environments, where they can discover effective policies directly from sensory inputs. However, these agents often exploit spurious correlations in the training data, resulting in brittle behaviours that fail to generalize to new or slightly modified environments. To address this, we introduce the Causal Object-centric Model Extraction Tool (COMET), a novel algorithm designed to learn the exact interpretable causal world models (CWMs). COMET first extracts object-centric state descriptions from observations and identifies the environment's internal states related to the depicted objects' properties. Using symbolic regression, it models object-centric transitions and derives causal relationships governing object dynamics. COMET further incorporates large language models (LLMs) for semantic inference, annotating causal variables to enhance interpretability. By leveraging these capabilities, COMET constructs CWMs that align with the true causal structure of the environment, enabling agents to focus on task-relevant features. The extracted CWMs mitigate the danger of shortcuts, permitting the development of RL systems capable of better planning and decision-making across dynamic scenarios. Our results, validated in Atari environments such as Pong and Freeway, demonstrate the accuracy and robustness of COMET, highlighting its potential to bridge the gap between object-centric reasoning and causal inference in reinforcement learning.