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

2026-05-18

AI Summary

Dense video generation in embodied control is computationally expensive and redundant, hindering efficient prediction of task-relevant visual states. This work proposes SWEET, a novel framework that leverages an image editing model as a sparse visual world model to generate task-specific keyframe sequences in a single stage, guided by language instructions and arrow-based spatial cues, and directly maps them to executable action chunks. Integrating the FLUX-Kontext editing model, a goal-conditioned diffusion-based action predictor, and a hybrid training strategy, SWEET substantially reduces inference costs while improving keyframe fidelity. Experiments demonstrate that SWEET accurately predicts critical keyframes on both DROID and RoboMimic datasets and enables end-to-end closed-loop execution from planning to robotic control.

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.

Problem

• sparse visual prediction
• embodied task execution
• image editing
• robot manipulation
• visual world modeling

Innovation

• image editing
• sparse world modeling
• embodied task execution
• visual prediction
• diffusion action prediction