Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Better Decisions through the Right Causal World Model

Elisabeth Dillies, Quentin Delfosse, Jannis Blüml, Raban Emunds, Florian Peter Busch, Kristian Kersting

TL;DR

The paper tackles generalization gaps in deep RL caused by reliance on spurious correlations. It introduces COMET, a pipeline that combines object-centric perception, symbolic regression, and LLM-based semantic annotation to derive exact causal world models (CWMs) from observations and internal states. The approach yields CWMs for Pong and Freeway that reveal true object dynamics and their causal dependencies, enabling better planning and reduced reliance on shortcuts. This work advances the integration of object-centric reasoning and causal inference in RL and suggests practical steps to leverage hidden internal states and LLMs for interpretable, generalizable agents.

Abstract

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.

Better Decisions through the Right Causal World Model

TL;DR

The paper tackles generalization gaps in deep RL caused by reliance on spurious correlations. It introduces COMET, a pipeline that combines object-centric perception, symbolic regression, and LLM-based semantic annotation to derive exact causal world models (CWMs) from observations and internal states. The approach yields CWMs for Pong and Freeway that reveal true object dynamics and their causal dependencies, enabling better planning and reduced reliance on shortcuts. This work advances the integration of object-centric reasoning and causal inference in RL and suggests practical steps to leverage hidden internal states and LLMs for interpretable, generalizable agents.

Abstract

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.

Paper Structure

This paper contains 4 sections, 2 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Method
  3. COMETs extracted world models
  4. Discussion and future work

Figures (2)

  • Figure 1: Deep RL agents learn undetectable shortcut. A deep (PPO) agent reaches the maximum score and selects the correct action in the depicted state of the training environment. The explanation map further leads external reviewers to think that the agent "understands" that it should return the ball behind the enemy (Left). In the test environment, where the enemy is hidden, the agent decides to go down, preventing it from catching the ball (Right). This illustrates that the agents learned to rely on the enemy's vertical position as an estimation of the ball's position.
  • Figure 2: Example of world models extracted by COMET. Top: The extracted world model for the Pong environment, that consist of the player, the enemy, their scores, and the Ball. Bottom: Part of the Freeway world model ($4$ depicted cars out of $10$). All the internal RAM variables extracted directly linked to the objects' properties (light blue) are on par with the annotated RAM provided by anand2019unsupervised. The variables used to update these properties (dark blue) are valid. For example, ram[58] corresponds to the speed of the Ball on the x-axis.