ID-WM is a hierarchical world modeling and planning framework that decouples temporally extended subgoal generation from local reachability. By treating diffusion noise as implicit actions, it constrains the planning search space to a unit Gaussian, enabling reliable subgoal generation. ID-WM outperforms the evaluated baselines and demonstrates vision-based subgoal generation.
Experiments
The experiments address four questions:
- Subgoal generation: Can ID-WM reliably generate diverse and feasible reachable states?
- Generalization: Can ID-WM solve long-horizon planning tasks to reach arbitrary, zero-shot goals at test time?
- Scaling: How does ID-WM perform across environments of increasing complexity?
- Visual planning: Is ID-WM effective for long-horizon planning directly from high-dimensional visual observations?
Environments
ID-WM is evaluated on three navigation and manipulation environment suites: Wall, PushT, and Maze. Task horizons range from hundreds to thousands of environment steps.
Wall
Wall is a simple 2D navigation environment consisting of two rooms separated by a wall with a door. The agent must navigate from a random starting location on one side of the wall to a random location on the other, passing through the door. The dataset consists of 1,920 randomly generated trajectories, each 50 steps long.
Baselines
ID-WM is compared with offline learning methods: goal-conditioned behavior cloning (GCBC), offline reinforcement learning (HIQL), diffusion planners (Diffuser and Hierarchical Diffuser), and planning with a single-level world model and model predictive control (MPC). Qualitative vision-based experiments also compare ID-WM with AVDC.
Single-Level World Model and Planning
This baseline employs a flat world model to perform single-level planning with the Cross-Entropy Method (CEM), without hierarchical abstractions.
Goal-Conditioned Behavior Cloning
GCBC is an imitation learning paradigm in which a policy maps the current state and a target goal observation directly to the expert action that leads toward that goal. The implementation extends Diffusion Policy by conditioning the policy on a future frame randomly sampled from the trajectory.
HIQL
HIQL is an offline reinforcement learning method that learns hierarchical representations by combining a high-level subgoal-as-action policy with a low-level subgoal-conditioned policy for long-horizon tasks.
Diffuser
Diffuser is a trajectory-level diffusion model that plans by sampling entire state-action sequences. The method replans after each action is executed in the environment.
Hierarchical Diffuser
Hierarchical Diffuser decomposes planning into two temporal scales. A high-level module synthesizes intermediate subgoals, which a low-level denoising network uses to generate a refined trajectory. Both levels are replanned after each action is executed in the environment.
AVDC
AVDC uses a video generation model to predict future states and inverse kinematics (IK) to derive the required actions. It is used to compare generated video sequences with ID-WM’s high-level subgoals.
HWM
HWM is a concurrent hierarchical latent world model that increases the high-level model’s frameskip and plans over learned macro-actions. The comparison uses the DINO-WM variant of HWM. Because the official implementation of this variant is not open-source, it was reimplemented using the reported settings and DINO-WM for unspecified details.
Generating Neighboring States with ID-WM
The experiments examine how accurately each level of ID-WM captures the underlying data distribution and predicts future states across different temporal granularities. For the high-level world model, the evaluation tests whether sampling from the latent prior yields diverse and reachable future states.
Scaling with Task Complexity
The experiments investigate how ID-WM scales with increasing environment size and goal horizon.
Scaling with Goal Horizon
Building on the PushT planning benchmark, which evaluates goals 25 steps away, the evaluation increases complexity by sampling validation-set goals that are an unspecified number of steps away.
Data Augmentation
- Rotation (30%): All positions and the T-block angle are rotated by an unspecified angle about the scene center. Velocities and actions are rotated accordingly.
- Displacement + Rotation (30%): A half-range displacement is combined with an unspecified small rotation.
Compute Resources
State-based world models are trained using 1 to 8 NVIDIA H200 Tensor Core GPUs. Distributed Data Parallel (DDP) is used to accelerate training when scaling to 8 GPUs. Baseline models are trained to convergence on a single NVIDIA H200 Tensor Core GPU, and planning evaluations also use a single NVIDIA H200 Tensor Core GPU.
For vision-based experiments, a CDiT backbone is trained using DDP across a cluster of 32 NVIDIA H200 Tensor Core GPUs.
High-Level World Model Training Convergence
At the beginning of training, predicted distributions are highly dispersed and often fail to preserve the structure of the PushT scene. Early samples do not consistently respect the object’s local geometry or the physically plausible displacement caused by agent-object interaction.
As training progresses, generated samples become more concentrated around coherent future configurations. By approximately 200 to 600 training steps, the model begins to capture the coarse reachable region for each temporal horizon. Later checkpoints sharpen the predicted distribution while preserving meaningful sample diversity.
Conclusion
ID-WM is a hierarchical world modeling and planning framework that decouples temporally extended subgoal generation from local reachability. Treating diffusion noise as implicit actions constrains the planning search space to a unit Gaussian, ensuring reliable subgoal generation. ID-WM outperforms the evaluated state-of-the-art baselines and demonstrates high-quality vision-based subgoal generation. Grounding hierarchical abstraction in explicit local dynamics is a step toward robust, long-horizon decision-making.
Limitations
ID-WM’s search-based optimization in a diffusion latent space incurs higher inference times than reactive policies or deterministic world models. Future work will explore amortized or gradient-based planning to address this bottleneck.
Although the hierarchy prevents planning complexity from scaling with trajectory length, complexity still scales with state-space dimensionality. Integrating compact, structured representation learning could further streamline planning in highly complex environments.
Acknowledgments and Disclosure of Funding
The authors thank Amir Bar, Alfredo Canziani, Yilun Kuang, Hengkai Pan, Mike Rabbat, Ying Wang, and Wancong Zhang for their valuable discussion and feedback.
This work was supported by AFOSR under grant FA95502310139, and by the Institute of Information & Communications Technology Planning & Evaluation (IITP) with a grant funded by the Ministry of Science and ICT (MSIT) of the Republic of Korea in connection with the Global AI Frontier Lab International Collaborative Research Agreement (No. RS-2024-00469482 & RS-2024-00509279).
Frequently Asked Questions
What is ID-WM designed to do? It combines hierarchical world modeling and planning, separating extended subgoal generation from local reachability.
How does ID-WM constrain its planning search space? It treats diffusion noise as implicit actions, constraining the search space to a unit Gaussian.
Which environments are used in the evaluation? The evaluation uses Wall, PushT, and Maze navigation and manipulation environments.
What are the stated limitations of ID-WM? Search-based optimization increases inference time, and planning complexity still scales with state-space dimensionality.