Object-Centric World Models for Causality-Aware Reinforcement Learning

AI summary

Overview

• Research area: Model-based deep reinforcement learning (MBRL), object-centric representation learning, and Transformer-based world models.
• Technical level: Advanced. The paper assumes familiarity with reinforcement learning (actor-critic methods, GAE, discount factors), Transformers (attention, queries/keys/values, positional encoding), and slot-based autoencoders.
• Scope: The paper introduces STICA, a unified MBRL framework in which an object-centric Transformer world model, policy network, and value network all operate on per-object latent tokens, with causality-aware attention guiding decision-making.

What This Paper Is About

Most world models for reinforcement learning learn a single "holistic" representation of the entire scene, which makes it hard to capture the individual objects in an environment and the interactions between them. This is a problem for environments that are high-dimensional, non-stationary (targets or objects appear and disappear), and filled with many interacting objects. The authors propose STICA, which decomposes observations into object-level latent states, predicts each object's dynamics with a Transformer, and uses explicit token-level causal relationships to focus the policy and value networks on task-relevant objects.

Key Contributions

1. A high-performance object-centric world model. STICA is presented as the first MBRL agent that extracts object-centric representations directly from observations without random exploration, supervision, or pretraining. It is evaluated on the Safety Gym benchmark (first-person-view, non-stationary, object-rich 3D environments) and the Object-Centric Visual RL (OCVRL) benchmark.
2. Causality-aware decision-making. The policy and value networks estimate token-level dependencies through a causal graph matrix and per-latent causality scores, and use these to reweight attention. Visualizing attention weights yields human-interpretable causal influences that STICA exploits.
3. End-to-end learning of object-centric representations. By separating static background information into a dedicated learnable, time-independent latent state and removing it from the object-centric representation, the authors improve the autoencoder's feature extraction, which prior work had typically achieved through supervised learning or pretraining.
4. Consistent empirical gains. STICA is reported to outperform state-of-the-art baselines in both sample efficiency and final performance across the tested benchmarks.

Main Findings

• Causal attention is the largest single contributor. In ablation studies on PointButton1 and CarButton1, the authors report that "the introduction of causal attention yields a substantial performance increase" (STICA vs. STICA w/o CA). Exact ablation curves are shown in Figure 5 but numeric values are not given in the text.
• Object-centric representations alone help only marginally. The variant that removes causal attention, background removal, and the object-centric policy/value Transformers—effectively equivalent to TWM—shows only a marginal gain from the object-centric world model by itself.
• Object-centric policy and value Transformers matter. Replacing them with the original architectures (STICA w/o CA+BR vs. STICA w/o CA+BR+TP+TV) produces a noticeable performance gain.
• Background removal gives a modest improvement. Removing the static background latent state (STICA vs. STICA w/o BR, or STICA w/o CA vs. STICA w/o CA+BR) improves performance modestly.
• Safety Gym results. STICA's mean score is 11.90 versus DreamerV3 at 10.20, TWM at 8.39, TD-MPC2 at 3.73, and PPO at 3.15. Normalized mean: STICA 5.49, DreamerV3 4.06, TWM 3.83, TD-MPC2 1.15, PPO 1.00 (normalized against PPO). Per-task STICA scores: PointGoal1 13.63, PointGoal2 13.64, PointButton1 11.52, PointButton2 5.97, CarGoal1 17.09, CarGoal2 18.27, CarButton1 9.65, CarButton2 9.25. STICA is best on all tasks except PointGoal1, which the authors describe as a simpler task with fewer and less diverse objects where holistic representations may suffice; there DreamerV3 scored 19.13 versus STICA's 13.63.
• OCVRL results. STICA achieves the best success rates on all three tasks: Object Goal 0.737, Object Interaction 0.333, Object Reach 0.867. For comparison, TWM scored 0.727 / 0.080 / 0.772, DreamerV3 scored 0.677 / 0.156 / 0.697, and TD-MPC2 scored 0.590 on Object Reach (it was omitted from the 2D tasks because it is designed for continuous action spaces).
• Object Interaction is the standout. The authors report a "drastically better" result on Object Interaction, where the agent directly interacts with objects, and attribute this to the object-centric world model and object-centric causal value network acting synergistically.
• Attention visualization confirms selectivity. The value network in STICA attends exclusively to reward-yielding target objects (the green button or green cylinder area), with no appreciable attention to other objects. The policy network also focuses on the target but pays some attention to other objects, which the authors say is necessary for action selection. In contrast, STICA w/o CA fails to focus on the target and pays attention to many objects, and STICA w/o BR attends heavily to the task-irrelevant background.
• Evaluation protocol. Safety Gym returns are averaged undiscounted returns over E = 10 episodes of T_ep = 1000 steps; training curves show averages over three runs with shaded standard deviations; OCVRL success rates are averaged over three runs.

Methodology in Plain English

STICA has three parts, trained in the standard model-based RL loop.

Object-centric world model. A slot-based autoencoder takes an image observation and produces n object slots through Slot Attention. Each slot is a 128-dimensional vector of logits defining 16 distinct categorical distributions of 8 categories, from which the latent state is sampled as 16 sets of 8-dimensional one-hot vectors. A separate learnable, time-independent latent state is dedicated to static background information such as environment layout and the agent's body. A spatial broadcast decoder reconstructs individual RGB images and unnormalized masks for each object plus a background image, and a softmax over the masks produces normalized masks used to composite a single reconstruction. The autoencoder is trained on reconstruction error, an entropy regularization term that prevents each latent from becoming deterministic, and a cross-entropy term aligning extracted latents with the dynamics model's predictions.

Dynamics model. A Transformer-XL aggregation model takes the history of rewards, object latent states, and actions as tokens and outputs hidden states for each latent plus an extra hidden state for the reward and discount. Causal masking prevents attention to future time steps, and the positional encoding depends only on time step, not on slot index, so the model is equivariant to slot ordering. Multilayer perceptrons then predict the next latent states (categorical), the reward (normal, via reparameterization), and the discount factor in [0, 1].

Causality-aware policy and value networks. Two separate Transformers take the latent state tokens plus an extra learnable token whose output gives the action or value estimate. A causal graph matrix encodes which token types can influence which others (target, causal object, non-causal object). An MLP estimates a causality score in [0, 1] for each latent state, representing the probability it is a causal object. These scores build a weight matrix, and the original attention weights are element-wise multiplied by the composition of the weight matrices with the causal graph, followed by row-wise L1 normalization—this is the causal attention used in every attention layer. Training uses A2C with advantages estimated by GAE.

Why This Matters

• Impact on research. The paper argues that holistic world models cannot capture object interactions, and shows that object-centric decomposition plus explicit token-level causal weighting is a practical route to better sample efficiency and final returns in object-rich, non-stationary, partially observable environments. It also removes the dependence of earlier object-centric MBRL methods on random exploration, supervision, or pretraining.
• Service robots that must identify which objects in a cluttered scene are relevant to the current instruction and ignore the rest.
• Autonomous driving, where the scene is high-dimensional, non-stationary, and composed of many interacting objects.
• Robotic manipulation, as suggested by the Object Interaction and Object Reach tasks, where the agent must push or touch a specific target object.
• Safety-critical navigation, since the Safety Gym environments involve goals, obstacles, and buttons with dense and sparse rewards.
• Industry relevance. Sample efficiency matters when real-time operation and physical device failures make extensive real-world interaction costly, as the authors note. Interpretable attention over objects is also useful for debugging and validating learned policies in deployed systems.

Future Directions

• Scaling to more complex and larger object sets. The authors note that PointGoal1, with fewer and less diverse objects, was the one task where STICA did not show a significant advantage; how far the approach extends as object counts grow is not established.
• Causal inference versus token-level dependency. The paper explicitly states that "causality" here means token-level dependency, not causal inference in the formal sense. Extending the framework to proper causal reasoning is an open question.
• Reducing reliance on reconstruction. The design keeps a reconstruction-based autoencoder; whether the object-centric and causal attention components transfer to reconstruction-free world models is not explored.
• Performance on safety costs. The paper uses only average undiscounted return and does not use the cost signals that Safety Gym was originally designed to measure, leaving safety-constrained object-centric decision-making as an unexplored direction.
• Compute and resource reporting. The content does not report compute budgets or training cost, so the practical cost of the method is not quantified.

Target Audience

Researchers and graduate students working on model-based reinforcement learning, Transformer-based world models, or object-centric representation learning. It is also relevant to practitioners who need sample-efficient agents for visually complex, object-rich, non-stationary environments and to those interested in interpretable attention over objects for policy analysis. Readers without a background in RL and attention mechanisms will find the method sections demanding.

Authors’ abstract

World models have been developed to support sample-efficient deep reinforcement learning agents. However, it remains challenging for world models to accurately replicate environments that are high-dimensional, non-stationary, and composed of multiple objects with rich interactions since most world models learn holistic representations of all environmental components. By contrast, humans perceive the environment by decomposing it into discrete objects, facilitating efficient decision-making. Motivated by this insight, we propose \emph{Slot Transformer Imagination with CAusality-aware reinforcement learning} (STICA), a unified framework in which object-centric Transformers serve as the world model and causality-aware policy and value networks. STICA represents each observation as a set of object-centric tokens, together with tokens for the agent action and the resulting reward, enabling the world model to predict token-level dynamics and interactions. The policy and value networks then estimate token-level cause--effect relations and use them in the attention layers, yielding causality-guided decision-making. Experiments on object-rich benchmarks demonstrate that STICA consistently outperforms state-of-the-art agents in both sample efficiency and final performance.