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

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.

• The paper introduces WorldDP, a hierarchical framework that combines object-centric world models with diffusion policies for robust multi-stage manipulation tasks.
• It leverages DINOv2 features with SAM2 segmentations and particle filtering to optimize subgoal planning and precise object-level state representation.
• Experimental results demonstrate significant improvements in task success for complex, multi-object scenarios compared to baseline approaches.

Hierarchical Object-Centric World Models and Diffusion Policies for Multi-Stage Robotic Manipulation

Overview and Motivation

The paper "Unifying Object-Centric World Models and Diffusion Policy: A Hierarchical Framework for Multi-Stage Robotic Tasks" (2606.08775) presents WorldDP, a new hierarchical control framework for robotic manipulation premised on integrating object-centric visual world models with goal-conditioned diffusion policies in a two-tiered architecture. The motivation stems from the critical limitations of contemporary visual world models — while effective for single-stage control, they are ill-suited to long-horizon, sequential multi-stage manipulation due to challenges in planning, subgoal generation, and latent representation of environmental dynamics at the object level.

WorldDP addresses these deficits by combining:

• Object-centric state representations built on top of DINOv2 patch features, guided by SAM2 segmentations;
• High-level subgoal planning using a learned world model capable of optimizing over latent action sequences via a particle filter;
• Low-level goal-conditioned diffusion policy for robust and efficient execution of planned subgoals.

This decomposition enables sequential, object-level reasoning while offloading fine-grained control to a diffusion policy, yielding marked improvements in multi-stage and composite manipulation tasks.

Figure 1: Overview of WorldDP: (a) the object-centric world model is used within a particle filter for subgoal optimization, and (b) a goal-conditioned diffusion policy tracks and executes these subgoals sequentially.

Object-Centric Representation and Encoder Design

Classical latent world models use patch-level encodings, which insufficiently capture fine-grained object dynamics, especially in settings with multiple interacting entities. WorldDP leverages an object-centric encoder (OCE) that utilizes DINOv2 patch features as the input, grouping them into slot-based representations corresponding to discrete environment entities such as robot, objects, and background.

Training leverages SAM2 to provide segmentation masks, supervising the slot assignment through a composite reconstruction and segmentation loss that combines $\ell_2$-reconstruction of patch features and a Tversky loss on mask segmentation.

Figure 2: Object-Centric Encoder pipeline with DINOv2 patch features, slot attention, SAM2 supervision, and dual loss objectives.

This design enables accurate slot assignment and more robust tracking of small objects across high-variance manipulation scenes, crucial for both downstream planning and control.

Hierarchical World Model and Planning via Particle Filtering

The world model is a conditional diffusion transformer (CDiT) mapping object-centric state-action sequences to future object-centric states. During planning, the model rolls out latent future trajectories by sampling action sequences, leveraging a particle filter (PF) for multi-modal optimization — a critical enhancement over unimodal CEM-based planning under the multi-modality of robotic tasks.

The cost function guiding this optimization operates directly over object embedding distances (MSE) augmented by a contact prediction cost, which signals critical manipulation phases (e.g., grasp events) and improves the temporal localization of subgoals. This explicit object- and contact-based cost is a principal source of WorldDP's planning improvements.

Goal-Conditioned Diffusion Policy for Local Execution

At the execution tier, subgoals output by the world model are realized by a goal-conditioned diffusion policy (DP), implemented as a transformer-based action generator. Conditioning occurs on the flattened object-centric current and goal state representations, as well as robot proprioception. This decouples long-horizon sequential planning from short-horizon robust control, capitalizing on the strengths of diffusion policies: sample-efficient, model-free execution that is tolerant to small errors in subgoal selection.

Experimental Evaluation

WorldDP is evaluated on OGBench-based manipulation environments spanning Cube-Single, Cube-Triple, Scene-Single-Direct, and Scene-Single-Composite. Each introduces increasing complexity in terms of number of objects, degrees of freedom, and compositional/casual dependencies.

Figure 3: Example start and goal images from the evaluated multi-stage robotic benchmarks.

Quantitative Results

WorldDP demonstrates strong empirical performance over baselines (DINO-WM, LeWM, HECRL, and vanilla DP) in both hard (Cube-Triple, Scene-Single-Composite) and simple (Cube-Single, Scene-Single-Direct) tasks, with particular advantage seen in fully sequential or causally entangled manipulations.

Strong numerical results include:

• Cube-Triple (3 cubes): WorldDP achieves $30\%$ full task success (vs. $12\%$ for HECRL$^*$ and $4\%$ for DP100).
• Scene-Single-Composite (Composite): $20\%$ completion rate, outperforming HECRL$^*$ and DP100 by a significant margin.
• Substantial reduction in error propagation and task failures versus flat or single-level planning architectures.

Component and Ablation Analyses

Ablations confirm that:

• Removing the OCE reduces task success substantially, with especially steep drops in multi-object scenarios.
• Using image-embedding costs (rather than object-centric) in planning degrades performance, especially for hard compositional tasks.
• Omission of the diffusion policy or contact predictor similarly reduces overall task completion.

  Figure 4: Success rates fall markedly when either the low-level diffusion policy or object-centric encoder is omitted.

Qualitative Insights

Figure 5: Scene-Single-Composite execution: WorldDP hierarchically decomposes the task into button and window/drawer phases, planning and executing subgoals sequentially via DP.

Visualizations of object-centric masks indicate clear segmentation and stable slot tracking even for small or partially occluded entities. Open-loop rollouts of the world model reveal accurate long-horizon predictions in latent space, decoded for visualization, further validating the learned dynamics.

Figure 6: Open-loop trajectory rollouts decoded from latent object-centric states, showing accurate temporal prediction over 4-second horizons.

Theoretical and Practical Implications

WorldDP provides a unified perspective on hierarchical control, object-centric abstraction, and generative policies within the context of vision-based robotics. Its main conceptual contributions include:

• Hierarchical Decomposition: Demonstrates the advantage of separating long-horizon planning (world model) from local execution (diffusion policy), enabling robust compositionality needed for multi-stage tasks.
• Object-Centric Planning: Reinforces the value of entity-level state abstraction for both representation and downstream planning optimizations, as evidenced quantitatively and qualitatively.
• Particle Filtering in Planning: Introduces PFs in latent trajectory optimization, yielding improved scaling to multi-modal settings inherent in complex robotics.

Practically, WorldDP extends the frontier for goal-directed robotic manipulation in tasks previously infeasible for flat or single-stage world models. However, latency introduced by planning (up to two minutes per task execution) highlights a need for further optimization (e.g., parallelized or gradient-based approaches), and the reliance on image-goal specification may limit application breadth.

Conclusion

WorldDP establishes a new paradigm for vision-based multi-stage robotic control, integrally connecting object-centric world modeling, hierarchical MPC-inspired planning, and diffusion-based local execution. Experimental results demonstrate consistent outperformance on composite, multi-object manipulation benchmarks, substantiated by strong ablation and qualitative visualizations. The architecture's layered compositionality and object-level reasoning broaden theoretical and practical understanding of scalable, generalizable robotic planning. Future work is likely to explore language-driven goal specifications, efficient planning, and the extension of these techniques to more open-ended, cross-domain manipulation scenarios.