SWEET: Sparse World Modeling with Image Editing for Embodied Task Execution

TL;DR

SWEET introduces a sparse visual planning framework using image editing to generate task-relevant keyframes for robot manipulation, reducing computational costs and improving reliability over dense video generation methods.

Contribution

The paper demonstrates that image editing models outperform video generation models for task-level state prediction and proposes SWEET, a novel sparse planning approach based on sequential image editing.

Findings

Image editing produces more reliable task keyframes than video generation.

SWEET improves keyframe prediction in unseen scenes.

The full pipeline enables robot actions from sequential keyframes.

Abstract

Visual prediction has emerged as a promising paradigm for embodied control, where future observations are generated and then translated into actions. However, dense video generation is computationally expensive and often unnecessary for many manipulation tasks, whose progress can be summarized by a small number of task-relevant visual states. In this work, we study whether image editing models can serve as sparse visual world models for robot manipulation by predicting task-level future states without dense video rollout. We first conduct a controlled comparison between the video generation model Wan2.2 and the image editing model FLUX-Kontext under the same robotic data setting, and find that image editing produces more reliable task-level keyframes with better visual fidelity and substantially lower inference cost. Motivated by this observation, we propose SWEET, a one-shot sparse visual planning framework that progressively generates a sequence of task-relevant manipulation keyframes through successive image editing, conditioned on language instructions and optional arrow-based spatial guidance. A goal-conditioned diffusion action predictor then converts adjacent imagined keyframes into executable action chunks. To reduce the mismatch between real and edited visual subgoals, we further introduce a mixed-training strategy with filtered edited targets. Experiments on DROID and RoboMimic show that SWEET improves keyframe prediction across seen and unseen scenes and enables a full pipeline from sequential keyframe planning to executable robot actions, suggesting that image editing is a promising and underexplored direction for embodied visual prediction.