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

2026-06-07

AI Summary

Existing visual world models struggle with multi-stage robotic manipulation tasks requiring complex temporal planning, as they are typically limited to single-stage scenarios. This work proposes WorldDP, a novel framework that, for the first time, integrates an object-centric world model with a diffusion policy within a hierarchical architecture: a high-level model dynamically generates feasible subgoals at runtime, while a low-level diffusion policy executes the corresponding subtasks. By decoupling environmental entities, the approach enables object-oriented sequential planning and incorporates model predictive control (MPC). Experiments demonstrate that WorldDP significantly outperforms current methods across multiple robotic benchmark tasks, validating the effectiveness of coupling physically aware planning with efficient execution strategies in multi-stage settings.

Abstract

Visual world models have shown great potential in learning complex system dynamics. Recent advancements leverage these models as transition functions within Model Predictive Control (MPC) frameworks to solve various control tasks. When applied to robotics, however, they are limited to single-stage tasks such as reaching or grasping, and struggle with multi-stage ones that demand complex sequential planning. In this work, we introduce WorldDP, a world model framework designed for multi-stage robotic manipulation. Our hierarchical approach utilizes a high-level world model as a transition function to optimize for feasible subgoals during runtime, which are subsequently reached by a low-level Diffusion Policy. To further aid in learning dynamics and planning, we incorporate object-centric representations that decouple environmental entities and enable us to plan sequentially with respect to each. Evaluated across several robotics benchmarks, WorldDP consistently outperforms existing baselines, validating that coupling the world model's physically grounded planning with diffusion policy's efficient execution yields superior multi-stage performance.

Problem

• multi-stage robotic tasks
• visual world models
• sequential planning
• object-centric representations
• robotic manipulation

Innovation

• object-centric world model
• diffusion policy
• hierarchical planning
• multi-stage robotic tasks
• model predictive control