TLDR: This research introduces DLPWM, an unsupervised object-centric world model that learns disentangled object representations directly from pixels, showing strong visual modeling capabilities. However, when applied to model-based control, policies trained on DLPWM underperform compared to DreamerV3 due to “representation shift” during multi-object interactions, particularly during physical contact. The paper highlights that while object-centric perception is robust for visual modeling, stable control requires addressing latent drift.
The quest to develop artificial intelligence agents that can reason about the world in a structured, human-like manner is a persistent challenge. A promising path towards this goal involves Object-Centric World Models (OCWMs), which aim to break down complex visual scenes into individual object-level representations. These structured abstractions have shown significant potential for improving how AI systems generalize and learn efficiently in various downstream tasks.
Seminal works in this area have demonstrated that object-centric representations, often learned without direct supervision, can facilitate robust object property prediction and exhibit resilience to certain out-of-distribution visual variations. Furthermore, recent studies suggest that these representations can lead to better compositional generalization with reduced computational overhead.
Building on these findings, a key question arises: can structured, object-centric representations improve policy learning in reinforcement learning (RL)? Although recent studies show gains in relational or compositional tasks, these often fail to generalize beyond the training distribution. The authors of this paper posit that this limitation arises from a misalignment between representation and control, where object-centric models still entangle task-relevant and irrelevant factors—yielding policies that are structured but brittle.
The core hypothesis of this research is that the key to unlocking the full potential of object-centric representations for policy generalization lies in the disentanglement of task-relevant features. By isolating the factors of variation that are crucial for a given task, the resulting latent space becomes a more reliable and robust foundation for the policy, particularly in novel or out-of-distribution scenarios. A policy that can consistently access and reason about these disentangled features should, in principle, generalize more effectively to unseen combinations of object properties and environmental conditions.
Disentangled object-centric methods pursue this goal by producing per-object latents that separate spatial attributes (e.g., position, scale) from appearance and dynamics. For example, Deep Latent Particles (DLP) represent objects as learned particles with explicit spatial and feature components, enabling interpretable, unsupervised learning from pixels. Its temporal extension, Deep Dynamic Latent Particles (DDLP), augments DLP with dynamics-aware particles for object-level video prediction.
To test their hypothesis, the researchers introduced DLPWM, a fully unsupervised, disentangled object-centric world model that learns object-level latents directly from pixels. In initial evaluations, DLPWM demonstrated strong performance in visual modeling, achieving high-fidelity reconstruction and accurate prediction, even in the presence of several out-of-distribution visual variations. This suggests that the model successfully learns a robust and generalizable representation of the visual scene.
However, when these learned representations are leveraged for downstream model-based control, policies trained on DLPWM latents underperform in comparison to the state-of-the-art holistic world model, DreamerV3. Through in-depth latent-trajectory analyses, the researchers identified a significant “representation shift” during multi-object interactions as a primary driver of this unstable policy learning.
The experiments benchmarked DLPWM against DreamerV3 on a standard block-lifting task and a generalized object-lifting task called the Generalization Arena, which allows control over object shapes and colors. DLPWM generally achieved better reconstruction performance, while prediction accuracy remained comparable between the two models.
When it came to policy learning, after 50,000 steps of offline DLPWM training, online policy learning was initiated. Both GNN-based and Transformer-based particle aggregators were tested. The results showed that both policies successfully learned to reach the target object (within the initial 100,000 steps) but struggled with grasping and lifting. The researchers hypothesize that this limitation arises from direct physical interactions between entities—when objects make contact, the particle representation may fail to preserve clear boundaries between them.
Further analysis of latent trajectories during contact events, as shown in the paper, revealed that near contact frames, the positional, scale, and visual feature components of the latent space exhibited noticeable changes. While such alterations should occur during physical interaction, the arm’s proximity alone could perturb the object’s latent features even before contact. This interference is hypothesized to contribute to the observed policy underperformance.
In conclusion, DLPWM is a disentangled object-centric world model that achieves strong reconstruction and prediction capabilities. However, its imagined latents sometimes lead to unstable policy learning, particularly during interaction-heavy task phases. The researchers hypothesize that brief, contact-induced perturbations and occasional slot-identity drift corrupt the per-slot inputs the policy uses.
As next steps, the authors propose to supply the policy with an exponential moving average (EMA) of each slot—concatenated with the raw slot (and optionally the slot-delta)—which aims to attenuate high-frequency corruption, is trivial to compute online in both real and imagined rollouts, and preserves fast signals when raw and delta are also provided.
Also Read:
- Multimodal Diffusion Forcing: A Unified AI Framework for Robust Robot Manipulation
- Navigating Uncertainty: A New Approach to Online Robust Reinforcement Learning
You can find the full research paper here: WHEN OBJECT-CENTRIC WORLD MODELS MEET POLICY LEARNING: FROM PIXELS TO POLICIES, AND WHERE IT BREAKS.
