World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models

Abstract

Building generalizable agents for diverse applications remains a fundamental challenge. While imitation learning-based policies succeed in specific training environments, they often fail to generalize to novel scenes and tasks. In this work, we propose World Action Planner, a robot planning system that leverages the reasoning capabilities of Vision-Language Models (VLMs) and the physical grounding of a multi-task pose-image conditioned world model. Our system enables an agent to propose initial action plans and iteratively refine them via optimization and search, reasoning over imagined world model rollouts. We demonstrate that our approach achieves superior performance across compositional tasks, new layouts, and zero-shot generalization scenarios, significantly outperforming state-of-the-art end-to-end policy models such as VLAs and WAMs.

We introduce World Action Planner, a planning system in which the agent reasons about and refines its action plans via world model imagination, generalizing to novel scenarios and solving new tasks at test time.

Algorithm

Generalizable Decision-Making

Compositional Generalization
We evaluate compositional generalization on LIBERO-Long tasks, each composed of two LIBERO-90 tasks seen during training. After completing the first subtask, the VLM proposes transition actions toward the next object, and the action-conditioned world model imagines the resulting rollout to verify that the policy can subsequently resume and complete the next subtask.

New Layout Generalization
We evaluate our method on LIBERO-Object tasks with modified object layouts. The VLM agent identifies the relocated target object and navigates the gripper to its proximity, where a weak policy can reliably complete the grasp. The action-conditioned world model then imagines placement trajectory, allowing the planner to detect and avoid potential collision risks before execution.

Zero-Shot Generalization
We evaluate on Robosuite StackCube and PickPlaceCan without expert demonstrations or a learned robot policy. Using only hard-coded grasp/release primitives, the planner leverages world model imagination to iteratively refine fine-grained actions, such as precise cube alignment for successful stacking.

Robot World Model

Pose-Image Conditioned World Model
We compute robot joint poses via forward dynamics and render them as joint skeleton images from the corresponding viewpoints, which are then used as conditioning inputs for action-conditioned world modeling.