Skip to content

Object-Centric World Models for Causality-Aware Reinforcement Learning

🕒 Published (v1): 2025-11-18 08:53 UTC · Source: Arxiv · Venue: AAAI · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

STICA (Slot Transformer Imagination with CAusality-aware RL) is a model-based RL agent that uses a slot-based autoencoder and Transformer dynamics model to build object-centric world models, then exploits learned per-object causal scores to bias attention in policy and value networks toward task-relevant objects. It outperforms DreamerV3, TWM, and TD-MPC2 on both Safety Gym and OCVRL benchmarks in sample efficiency and final performance.

Problem

Existing world models (Dreamer series, STORM, TWM) learn holistic environment representations that fail to disentangle individual object dynamics and interactions, degrading performance in high-dimensional, non-stationary, multi-object environments. Prior object-centric MBRL methods require random exploration, supervised pretraining, or external datasets, limiting generality.

Method

STICA has three components:

  1. Slot-based AutoEncoder: Slot Attention extracts \(n\) categorical latent slots \((z_t^1,\ldots,z_t^n)\) per timestep. A learnable time-independent token \(z_{BG}\) absorbs static background, freeing object slots to capture only dynamic entities. A spatial broadcast decoder reconstructs per-object RGB images and masks, which are softmax-blended into the full observation.

  2. Transformer-XL Dynamics Model: Takes object slots, action, and reward as separate tokens; predicts next-slot distributions (categorical), reward (Gaussian), and discount (Bernoulli) via per-token MLPs. Cross-entropy alignment loss couples the autoencoder and dynamics model end-to-end without pretraining.

  3. Causal Policy and Value Networks: Each uses a Transformer that replaces standard self-attention with causal attention. An MLP estimates a per-slot causality score \(p_t^k \in [0,1]\). A weight matrix \(W_t\) built from these scores is used to modulate attention weights via element-wise multiplication with a fixed causal graph \(G\), amplifying influence from causal objects and suppressing irrelevant ones. Policy learning uses A2C with GAE in imagination rollouts.

Key Contributions

  • First MBRL agent achieving end-to-end object-centric representation learning without random exploration, supervision, or pretraining.
  • Background removal via a time-independent latent \(z_{BG}\) that separates static scene structure from object dynamics.
  • Causal attention mechanism in both policy and value networks that learns to attend selectively to task-relevant objects via differentiable causality scores.
  • Unified framework where world model, policy, and value networks are all object-centric Transformers.

Results

Safety Gym (8 tasks; baselines: PPO, TWM, DreamerV3, TD-MPC2): - STICA achieves best normalized mean return (5.49) vs. DreamerV3 (4.06), TWM (3.83), TD-MPC2 (1.15), PPO (1.00). - Best on 7/8 tasks; marginally behind DreamerV3 only on PointGoal1 (simpler, fewer objects). - Highest absolute scores on CarGoal1 (17.09), CarGoal2 (18.27), CarButton1 (9.65).

OCVRL (3 tasks; baselines: TWM, DreamerV3, TD-MPC2 on 3D only): - Object Goal: STICA 0.737 vs. TWM 0.727, DreamerV3 0.677. - Object Interaction: STICA 0.333 vs. TWM 0.080, DreamerV3 0.156 (largest margin). - Object Reaching: STICA 0.867 vs. DreamerV3 0.697, TD-MPC2 0.590.

Ablation (PointButton1 & CarButton1): - Causal attention (CA) contributes the largest single gain. - Background removal (BR) adds a modest further improvement. - Object-centric Transformer policy/value heads (TP/TV) outperform MLP heads even without CA.

Limitations

  • "Causality" is token-level attention dependency, not interventional causality; the paper explicitly disclaims causal inference semantics.
  • Object-slot count \(n\) is fixed and must be set a priori; handling scenes with variable or unknown object counts is unaddressed.
  • Evaluated only on Safety Gym and OCVRL; no real-robot or pixel-diverse (e.g., Atari, DMControl) validation.
  • Performance advantage on PointGoal1 (simplest task) is absent, suggesting diminishing returns when holistic representations suffice.
  • Scalability to large \(n\) or very high-resolution inputs is not analyzed.

Relevance to Agentic AI / LLM Agents

Object-centric world models are a structural prior for agents operating in cluttered, multi-entity environments — a key challenge for embodied and robotic agents. STICA's causality-aware attention provides a differentiable mechanism for agents to identify which entities in the scene are action-relevant, directly addressing the perception bottleneck in complex agentic tasks. The end-to-end training without pretraining or supervision makes this approach more practical for deployment in novel environments. For LLM-based agents, the decompose-then-reason paradigm (slot → causal score → selective attention) mirrors structured reasoning over entity graphs, suggesting potential integration with tool-use or scene-graph representations in vision-language agents.