STORM: Search-Guided Generative World Models for Robotic Manipulation

TL;DR

This work tackles the challenge of long-horizon, physically grounded planning in robotic manipulation by moving beyond purely language-guided reasoning. It introduces STORM, a framework that combines a diffusion-based Vision-Language-Action policy with a reward-augmented generative video world model and Monte Carlo Tree Search to perform explicit lookahead over visually grounded futures. The approach achieves state-of-the-art performance on the SimplerEnv benchmark, with a 51.0% average success rate and substantial improvements in video-prediction fidelity (FVD reduction >75%). STORM also demonstrates robust failure recovery through re-planning, highlighting the practical benefits of search-guided generative world models for manipulation tasks.

Abstract

We present STORM (Search-Guided Generative World Models), a novel framework for spatio-temporal reasoning in robotic manipulation that unifies diffusion-based action generation, conditional video prediction, and search-based planning. Unlike prior Vision-Language-Action (VLA) models that rely on abstract latent dynamics or delegate reasoning to language components, STORM grounds planning in explicit visual rollouts, enabling interpretable and foresight-driven decision-making. A diffusion-based VLA policy proposes diverse candidate actions, a generative video world model simulates their visual and reward outcomes, and Monte Carlo Tree Search (MCTS) selectively refines plans through lookahead evaluation. Experiments on the SimplerEnv manipulation benchmark demonstrate that STORM achieves a new state-of-the-art average success rate of 51.0 percent, outperforming strong baselines such as CogACT. Reward-augmented video prediction substantially improves spatio-temporal fidelity and task relevance, reducing Frechet Video Distance by over 75 percent. Moreover, STORM exhibits robust re-planning and failure recovery behavior, highlighting the advantages of search-guided generative world models for long-horizon robotic manipulation.

Figures (5)

• Figure 1: The overall architecture of STORM. The decision loop follows Eq. \ref{['eq:storm_framework']}: MCTS orchestrates the process, using the VLA ($\pi_{\text{vla}}$) to propose candidate actions and the video predictor ($M_w$) to simulate their outcomes, ultimately selecting the optimal action $A_t^*$.
• Figure 2: Overview of the STORM framework, illustrating the decision loop orchestrated by MCTS. It integrates the VLA policy for proposing diverse action candidates and the video world model for simulating visual outcomes and rewards. The dashed loop represents iterative simulations (selection, expansion, evaluation, backpropagation) for foresight-driven planning. Note: STORM enables re-planning for failure recovery by grounding search in explicit visual rollouts.
• Figure 3: Qualitative results for video prediction (Task: Put Carrot on Plate). The prediction conditioned on the VLA's action (middle) aligns well with the ground truth (bottom).
• Figure 4: Case study on the "Put Carrot on Plate" task, demonstrating STORM's ability to recover from failure. Top: The baseline CogACT model fails, getting stuck in a repetitive loop after initial unsuccessful grasp attempts. Bottom: Our model, STORM, uses its lookahead planning to re-evaluate after the same initial failures and finds a new, successful trajectory to complete the task.
• Figure 5: Reward supervision is critical for learning a high-fidelity world model. The radar chart compares video prediction metrics for a model trained with our full objective ('action+reward') versus one without reward supervision ('action-only'). The reward-augmented model's superior performance across all axes demonstrates that this signal compels the model to learn task-relevant causal structures, moving beyond superficial visual patterns to enable effective, foresight-driven planning.