Unifying Object-Centric World Models and Diffusion Policy: A Hierarchical Framework for Multi-Stage Robotic Tasks

arXiv:2606.08775 · cs.RO, cs.AI · Submitted 2026-06-07 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: Robotics Radio. Generated commentary on the latest robotics and control papers.

Rosa: I'm Rosa, and with me are Dev and Taro, guest researcher.

Dev: Today's paper: "Unifying Object-Centric World Models and Diffusion Policy".

Rosa: Visual world models have shown great potential in learning complex system dynamics,

Dev: First, who's behind it and why it matters.

Paper summary: Rosa: Well, I'm really curious about this WorldDP paper. It sounds like they've tackled a real limitation in visual world models by moving beyond single-step tasks to handle those complex sequences that actually happen in the real world.

Dev: That’s what I was hoping to hear, Rosa. From an engineering standpoint, the idea of using a high-level model for subgoals and then having a low-level policy handle the execution sounds like it could actually manage the complexity of continuous action spaces better than just trying to plan everything at once.

Taro: I'm interested in how they plan those subgoals; when things go wrong in a multi-stage task, what does this framework do to recover? I want to know what happens when the world doesn't behave exactly as predicted by the model.

Rosa: Exactly, Taro. The paper introduces WorldDP as a hierarchical framework that uses an object-centric world model for high-level subgoal optimization and then employs a diffusion policy for the low-level execution of those subgoals during runtime. This structure is designed specifically to address the difficulty of planning long sequences in robotic manipulation tasks.

Dev: So, to put it plainly, they’re using this object-centric representation to decompose a big task into smaller goals that the diffusion policy can then pursue sequentially. That sounds like a way to manage the computational load while keeping control tight during execution.

Taro: And I see how decomposing it helps with planning for sequential control; if we can focus on achieving one critical state, like gripping an object, at a time, it might be easier to handle unexpected disturbances between those steps. But what about the accuracy of those initial subgoals?

Rosa: That’s a very valid point regarding the quality of those subgoals. The paper addresses this by including a contact predictor that signals when the robot actually interacts with an object, and this signal gets incorporated into the MPC cost function to encourage subgoals to capture those pivotal manipulation phases, like when you need to actually grasp something.

Dev: A contact predictor sounds crucial for grounding the planning in physical reality rather than just abstract state predictions. That means we're not just optimizing for a point in space, but for an action that results in actual contact. How does that affect the loop rate and latency?

Paper summary: Taro: The paper mentions they use a Particle Filter optimization instead of something like the Cross-Entropy Method because it handles the inherent multi-modality of robotic action spaces better, which suggests they're trying to find a broader set of possible good solutions rather than just one best guess.

Rosa: That particle filter approach is interesting because robotic tasks are naturally multimodal; there are often several different ways to achieve the same objective, so having multiple hypotheses helps explore those possibilities during the planning phase. We see this idea being used in their test-time world model optimization within a particle filter to find those optimal latent action sequences.

Dev: I worry about that complexity impacting the loop rate; if running a particle filter adds significant overhead, we have to make sure the overall latency doesn't slow down the real-time control loop too much, especially when dealing with continuous action spaces. We need to see how fast that planning step actually runs in practice.

Taro: If you look at the context vectors feeding into their diffusion policy—things like end-effector position and velocity—that suggests they're trying to give the low-level execution a good sense of where it needs to be relative to its current state, which should help it track those subgoals effectively.

Rosa: That’s right; the diffusion policy is conditioned on these object-centric representations flattened into a vector of shape N by d, along with those context vectors, which is what allows it to generate forty-step action sequences for executing the subgoals. It seems like a solid coupling between high-level planning and low-level execution.

Dev: So we have this hierarchy: the world model does the slower, more complex reasoning about objectives, and then the diffusion policy handles the faster, iterative steps to actually get there—that sounds like a good way to decouple the heavy lifting. But what if that contact predictor misfires?

Paper summary: Taro: That’s where robustness comes in; since they train this system on multi-stage tasks from benchmarks like OGBench, they are testing how well it handles the sequential nature of these tasks. The goal is to make the subgoals so well-defined that even if the execution drifts slightly, the diffusion policy can compensate for suboptimal subgoals to a good extent.

Rosa: It seems they believe that this unification of physically grounded planning with efficient execution yields superior performance on these long-horizon manipulation challenges compared to existing methods. We see evidence of this in their experimental validation across several benchmarks.

Dev: I'm glad they validated it on those benchmarks, but for me, the real test is the generalization outside the lab environment. How long can we expect this system to reliably operate without needing constant retraining when faced with entirely new objects or environments?

Taro: That’s a big question about real-world deployment. The success rate on tasks like Cube-Triple achieving one hundred percent is impressive, but deploying it means dealing with dynamic elements that aren't perfectly modeled beforehand. We need to see if the object-centric representation is flexible enough for novel entities and situations not seen in training.

Rosa: Indeed, the paper confirms that both the low-level Diffusion Policy and the Object-Centric Encoder are crucial components for achieving peak performance, which suggests that optimizing just one part of this system won't give you those strong results. The architecture seems to be built so that each layer contributes meaningfully to solving the overall problem.

Dev: So we have a framework where the object-centric world model decomposes the task, and the diffusion policy executes it robustly, but we still need to keep an eye on that latency when running these complex planning steps in a live loop. That seems like a challenging balance for an engineer to strike.

Taro: The implication here is that future autonomy research should focus on creating object-centric representations that are inherently robust to environmental noise and unexpected interactions, because the current success relies heavily on the model accurately predicting those critical contact moments.

Rosa: It really does point toward a future where complex robotic manipulation isn't just about learning one single move, but about mastering a sequence of carefully planned interactions guided by an object-centric understanding of what needs to happen next. This is what WorldDP offers for multi-stage tasks.

Conclusion: Rosa: So, we're wrapping up our discussion on WorldDP, which is titled "Unifying Object-Centric World Models and Diffusion Policy: A Hierarchical Framework for Multi-Stage Robotic Tasks."

Dev: That title really tells you a lot about what they’ve built—it's combining two different ways of looking at the world to handle those long sequences we talked about.

Taro: I think the core idea is that this framework structures complex manipulation by using a high-level model to figure out *what* needs to happen, and then a low-level policy handles the messy execution of each step.

Rosa: Exactly, and the authors are showing how they use object representations as the bridge between those two layers.

Dev: From an engineering standpoint, that unification suggests they managed to keep control tight during execution while still having a robust plan for the overall task structure.

Taro: It’s interesting to see how they tackle that inherent multi-modality of robotic actions using their specific planning mechanism within the world model.

Rosa: And what we're seeing here is a move toward systems that don't just learn one movement, but a coherent sequence of interactions guided by an object-centric understanding.

Dev: But I still have to wonder about how this translates when we take it out of the controlled lab environment and into something more unpredictable, like an open warehouse.

Taro: That’s the next big hurdle; if the system relies so heavily on predicting those specific contact moments accurately, how resilient is that structure when things go completely off-script?

Rosa: We'll have to dig into those experimental results to see how well they handle those real-world variances, and I want to know how long we can realistically expect this kind of autonomy in practice.

Tandon School of Engineering, New York University · Courant Institute of Mathematical Sciences, New York University · AMI Labs

cs.RO, cs.AI

Submitted: 2026-06-07

Updated: 2026-10-07

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 87/100

The gist: Visual world models have shown great potential in learning complex system dynamics, but they are currently limited to single-stage robotic tasks and struggle with multi-stage ones that demand complex

Key concepts

Object-Centric Encoder (OCE)
This component processes input images using DINOv2 features to create discrete vector embeddings for entities like the robot and objects. These embeddings form the core 'object-centric representations' used by the world model to understand the environment, focusing on what is physically present.
World Model Planning (MPC)
The world model acts as a high-level planner within a Model Predictive Control (MPC) loop. It uses this model to generate feasible subgoals during runtime. This planning involves simulating future states using dynamics predictions to determine the best sequence of actions needed to achieve the overall task goal.
Diffusion Policy (DP)
This is the low-level execution component, a goal-conditioned policy trained via diffusion methods. It takes object representations and context vectors as input to generate 40-step action sequences. This policy is responsible for the fine, robust control needed to reach the specific subgoals set by the world model.

Terminology

Summary

Visual world models have shown great potential in learning complex system dynamics, but they are currently limited to single-stage robotic tasks and struggle with multi-stage ones that demand complex sequential planning. This work introduces WorldDP, a novel hierarchical framework that unifies object-centric world models with diffusion policies to solve complex, multi-stage robotic manipulation tasks by using a high-level world model for subgoal optimization and a low-level diffusion policy for execution.

World Model Architecture and Object Representation

The framework begins by processing the input image through an Object-Centric Encoder (OCE) that leverages DINOv2 patch features. This encoder refines initial slots, representing environment entities such as the robot, objects, and background, into discrete vector embeddings called object-centric representations. The state representation for the world model is formed by aggregating these object-centric embeddings. The training of the OCE is guided by a dual objective: a reconstruction loss between image patches and reconstructed patch embeddings, and a mask segmentation loss against ground-truth masks generated using SAM2 to provide privileged guidance during training.

Hierarchical Planning via World Model

The world model serves as the high-level planner within an MPC framework. It is used to optimize for feasible subgoals during runtime. This planning process involves executing multi-step autoregressive rollouts through a Conditional Diffusion Transformer (CDiT) dynamics prediction model, which takes the object-centric states as input. To find optimal latent action sequences, the method employs a Particle Filter optimization rather than Cross-Entropy Method, as it better accommodates the inherent multi-modality of robotic action spaces. The cost function used during this planning phase is formulated to be object-centric, focusing on minimizing the Mean Squared Error relative to target object goal states and incorporating a contact prediction cost.

Low-Level Execution via Diffusion Policy

The subgoals generated by the world model are subsequently realized by a low-level, goal-conditioned Diffusion Policy (DP). This DP is adapted from existing transformer-based models but is conditioned on the object-centric representations flattened to a vector of shape N ×d, along with context vectors including robot end-effector position and velocity. The policy is trained to generate 40-step action sequences and operates autoregressively during inference, iteratively feeding back predicted states into the CDiT model. This approach allows the diffusion policy to be robust for the short-horizon control needed to reach subgoals and compensates for suboptimal subgoals effectively.

Crucial Enhancements for Multi-Stage Success

WorldDP incorporates specific mechanisms designed to improve sequential planning and subgoal quality. Key enhancements include:

  1. A contact predictor that signals when the robot interacts with an object, which is incorporated into the MPC cost function to encourage subgoals to capture critical manipulation phases, such as grasping objects.

  2. The use of a Particle Filter for planning, which is superior to methods relying on single Gaussian distributions because it maintains a diverse set of hypotheses (particles) to better capture these varied solutions.

  3. The object-centric cost function that isolates the N−2 objects from the predicted state, ensuring optimization is directly related to manipulating physical entities.

Experimental Validation and Results

Evaluations across several robotics benchmarks, including Cube-Single, Cube-Triple, Scene-Single-Direct, and Scene-Single-Composite tasks from the OGBench benchmark, consistently demonstrate that WorldDP consistently outperforms existing baselines. Specifically, for the Cube-Triple environment (3 cubes manipulation), WorldDP achieved a success rate of 100%, significantly outperforming methods like DinoWM (62%) and LeWM (74%). The framework's success is attributed to the unification of world model planning with diffusion policy execution, proving that coupling the world model’s physically grounded planning with diffusion policy’s efficient execution yields superior multi-stage performance. Furthermore, ablation studies confirm that both the low-level Diffusion Policy and the Object-Centric Encoder are crucial components for achieving peak performance.

Conclusion

WorldDP introduces a novel hierarchical framework that integrates object-centric world models with diffusion policies to solve complex, multi-stage robotic tasks. The method's success is rooted in its ability to decompose sequential tasks into structured planning phases—where the world model plans subgoals and the diffusion policy executes them—while utilizing object-centric representations and particle filter optimization to handle the inherent complexity and multimodality of robotic action spaces. This unification yields superior generalization for long-horizon manipulation.

The gist

WorldDP, a novel hierarchical framework that integrates object-centric world models with diffusion policies to solve complex, multi-stage robotic manipulation tasks by using a high-level world model for subgoal optimization and a low-level diffusion policy for execution. This work introduces WorldDP, a world model framework designed for multi-stage robotic manipulation.

Improvements for AI systems

Here are the specific improvements that can be made to existing AI systems by implementing the WorldDP framework, and what these improved systems will be capable of doing:

The implementation of WorldDP provides several transformative capabilities for robotic learning systems, primarily by enabling robust, long-horizon, multi-stage manipulation through a hierarchical planning structure.

Here are the specific improvements:

  1. A shift from single-stage task solvers (like standard Diffusion Policies or traditional MPC) to a unified framework capable of solving complex sequential problems.

  2. The integration of object-centric representations for state representation, moving beyond raw pixel features or simple patch embeddings.

  3. The use of a Particle Filter optimizer within the world model planning loop to handle multi-modal action spaces effectively during subgoal generation.

Here is what the improved AI system (WorldDP) can do:

  1. It can perform complex, multi-stage robotic tasks that require sequential subgoals, such as:

  2. Rearranging multiple distinct objects simultaneously (e.g., the Cube-Triple task), where success depends on achieving precise target configurations for every entity in the environment.

  3. Execute intricate operations involving causal dependencies and mechanism interactions, like those seen in Scene-Single-Composite, where prerequisite actions (e.g., pressing a button) must precede complex spatial manipulation (e.g., moving a drawer or window).

  4. Achieve superior performance on long-horizon manipulation tasks compared to single-stage methods, effectively decomposing a long sequence of required actions into manageable, temporally ordered subgoals that are robustly tracked by the low-level policy.

  5. Improve generalization across different task configurations (e.g., varying object counts or environmental layouts) by learning an object-centric state representation guided by segmentation masks (via SAM2), rather than relying solely on generic patch features from vision encoders like DINOv2.

  6. Generate more critical and precise subgoals (like the exact moment of grasping an object) by incorporating a contact predictor loss into the world model planning cost function, leading to better physical interaction success rates.

Sources

Related papers