Source-linked AI summary
ImageWAM: Do World Action Models Really Need Video Generation, or Just Image Editing?
Yuyang Zhang, Wenyao Zhang, Zekun Qi, He Zhang, Haitao Lin, Jingbo Zhang, Yao Mu, Xiaokang Yang, Wenjun Zeng, Xin Jin
TL;DR
Video-generation-based WAMs are costly and may model visual details or future trajectories that are weakly related to robot actions. ImageWAM instead repurposes pretrained image editing models for action prediction, achieving a 93.56% success rate on RoboTwin (Random) while matching state-of-the-art WAM performance.
Problem
Video-based WAMs incur costly inference and must predict action-irrelevant details and difficult multi-frame futures, motivating alternatives focused on instruction-guided visual change.
Method
ImageWAM reuses instruction-conditioned image-editing representations, including denoising transformer KV caches, to condition an action prediction expert without explicit future-video synthesis.
Results
93.56% success rate on RoboTwin (Random), substantially outperforming all VLA baselines and reaching performance comparable to state-of-the-art WAM models.
Takeaways & Limitations
Image editing provides a promising alternative backbone for world-action modeling and broadens visual generative pretraining beyond video-based modeling.
Takeaways & Limitations
Video-generation-based WAMs can produce visible artifacts around task-relevant objects in imagined future frames.
Abstract
from arXiv · showhide
World Action Models (WAMs) commonly rely on video generation to bridge visual world modeling and robot control. However, video-based WAMs face three coupled limitations: dense multi-frame future tokens make inference costly, full video prediction spends capacity on action-irrelevant temporal and appearance details, and long-horizon future imagination may introduce errors that mislead action prediction. These issues raise a simple question: Does world action model really need video generation? We propose ImageWAM, a simple WAM framework that repurposes pretrained image editing models for robot action prediction. In contrast to video generation, image editing provides a better-matched prior: it only needs to model a target-frame transformation, focuses on action-relevant current-to-target visual differences, and grounds task instructions to localized visual changes through edit pretraining. In practice, ImageWAM does not decode the target frame at inference time; instead, it conditions a flow-matching action expert on the KV caches produced by image-editing denoising, using them as a compact world-action context. ImageWAM outperforms standard VLA baselines and matching competitive WAMs without additional policy pretraining across different simulator and real-world experiments. It also reduces FLOPs to 1/6 and latency to 1/4 of video-based WAMs. Attention analysis further shows that editing caches focus on task-relevant change regions, supporting image editing as an effective alternative to video-based world-action modeling.
1 Introduction
The introduction argues that video-generation world-action models expend costly computation on dense, action-irrelevant future-video details. ImageWAM instead repurposes pretrained image-editing models to produce compact, instruction-conditioned representations for robot action prediction.
- Motivation: Video-generation backbones model complete future videos, including appearance, background, camera, and temporal details that may be weakly related to the next robot action.Generating many spatio-temporal tokens across multiple frames also makes inference costly for real-time robot control.
- Motivation: Image editing matches manipulation because it transforms the current scene according to language, emphasizing instruction-guided, task-relevant visual change rather than photorealistic future video.For many manipulation tasks, the essential signal is a transformation toward a desired visual state.
- Motivation: Image editing pretraining aligns language with what should change, where it should change, and how the instruction specifies that change.The introduction identifies instruction-to-change alignment as a key property supporting policy learning.
- ImageWAM: ImageWAM repurposes pretrained image-editing models as instruction-conditioned visual backbones, extracting editing-aware representations from the current observation and instruction for an action-prediction head.The framework does not aim to generate visually appealing edited images or use editing models as goal-image generators.
- Results: ImageWAM improves over standard visual and vision-language backbones under comparable action-prediction architectures, while analyses identify instruction conditioning and editing-oriented feature extraction as important.These findings indicate that the gains are not merely due to stronger image recognition or language alignment.
2 Related Works
Prior image-editing research has advanced from simple object edits to complex spatial, semantic, and knowledge-driven modifications, while this work reinterprets editing representations for robot policy learning. In parallel, world action models have used video generation as predictive visual planning followed by action decoding.
- Text-Guided Image Editing: Recent image-editing models extend text-guided edits from simple object changes to complex spatial, semantic, and knowledge-driven modifications.Prior work primarily evaluates perceptual quality and instruction fidelity.
- Text-Guided Image Editing: Image editing provides source-conditioned, change-centric representations that can serve as compact world-action backbones for robot policy learning.This study examines image editing from a robotics perspective rather than focusing only on perceptual quality and instruction fidelity.
- Video Generation for Robot Policy Learning: World action models use video generation as explicit visual planning, predicting complete future videos or visual rollouts from current observations and task context.The predicted rollout is translated into executable actions by an inverse dynamics model or action decoder.
3 Method
ImageWAM replaces dense future-video prediction with endpoint image editing and reuses intermediate editing KV caches as compact, task-conditioned context for action generation. It jointly trains the image-editing branch and flow-matching action expert, then uses a single editing forward step at inference.
- ImageWAM pipeline: ImageWAM predicts a task-relevant endpoint frame rather than a full future trajectory, providing a compact intermediate that summarizes the instruction-specified visual transformation.This avoids generating dense spatio-temporal tokens across multiple future frames while preserving reason-before-act policy learning.
- Editing-cache conditioning: The model reuses intermediate transformer key-value caches from image-editing denoising as conditioning context for the action expert instead of decoding the edited image.The caches encode task-conditioned visual transformation information after the visual latent interacts with the task instruction.
- Joint training objectives: During training, ImageWAM samples an editing timestep, predicts the endpoint image’s velocity field, and jointly optimizes image-editing and action flow-matching objectives.The action expert conditions on the current observation, task instruction, and editing cache; the total objective is L = Lact + Limg.
- Inference: At inference, ImageWAM performs one editing-branch forward step at a fixed timestep to obtain caches, then denoises action samples conditioned on those caches.It neither runs the full editing denoising trajectory nor decodes a complete edited image.
- Inference: ImageWAM uses a single set of layer-wise editing caches directly for action prediction, avoiding the dense future-video token instantiation required by standard video WAMs.A Fast-WAM-style comparison instead removes future video tokens only at test time and conditions actions on current-context KV caches.
4 Experiments
ImageWAM performs strongly across simulated and real-world manipulation benchmarks without additional embodied policy pretraining. Its image-editing caches also provide task-focused, artifact-avoiding action context with substantially lower inference cost than video-based WAMs.
- Evaluation Setup: ImageWAM uses only downstream benchmark demonstrations, evaluating LIBERO, LIBERO-Plus, RoboTwin 2.0, and real-world dual-arm manipulation tasks.The real-world setup uses four tasks spanning long-horizon, occluded, fine-grained, and deformable-object manipulation.
- Benchmark Results: 98.4% average success on LIBERO and 83.1% on LIBERO-Plus demonstrate competitive performance across standard and distribution-shifted simulation.ImageWAM remains competitive with video-generation WAMs and pretrained VLAs without data pretraining.
- Real-world Results: 84.5% average real-world success surpasses π0 (55.8%), π0.5 (72.3%), and FastWAM (79.0%) across all four evaluated tasks.Compared with FastWAM, gains are 6, 9, 1, and 6 points on T1 through T4, respectively.
- Analysis: ImageWAM attends to manipulated objects, target receptacles, and contact areas while suppressing backgrounds, and avoids future-video artifacts that can mislead action prediction.It uses image-editing caches directly as compact action-conditioning context without decoding future frames at inference time.
- Efficiency: 263 ms latency and 9.7 FLOPs, versus 1081 ms and 63.65, show substantially lower inference cost while maintaining competitive task success.ImageWAM obtains one image-editing cache set from a single editing-branch forward step rather than dense future-video tokens.
5 Conclusion
ImageWAM uses image editing instead of video generation as its WAM backbone, predicting a single future frame to provide intermediate representations for action modeling and end-to-end policy learning. It achieves a 93.56% success rate on RoboTwin (Random), substantially outperforming all VLA baselines.
- Conclusion: ImageWAM employs an image-editing model rather than a video-generation model as its WAM backbone.The paper motivates image editing as a general task requiring both visual understanding and generation.
- Conclusion: Predicting a single future frame provides strong intermediate representations for the action model and enables end-to-end policy learning.
- Conclusion: 93.56% success rate on RoboTwin (Random) substantially outperforms all VLA baselines.
Appendix
The appendix details ImageWAM’s multimodal architectures, action-head initialization, and model-specific backbones across OmniGen2, FLUX.2, and Ovis-U1. It also specifies shared training settings and dataset-specific preprocessing, prediction horizons, and schedules.
- Architectures: OmniGen2, FLUX.2[klein], and Ovis-U1 all use a Mixture-of-Transformers structure for multimodal joint modeling.
- Architectures: The OmniGen2 variant combines a pretrained Qwen2.5-VL-3B LLM with OmniGen2’s DiT over language, reference-image, and future-noisy-frame tokens.
- Action-head initialization: Action DiT initialization copies and interpolates image-editing weights, adds action projections, and prevents early noisy action gradients from affecting the visual model.
- Architectures: The FLUX.2 variant uses pretrained Qwen3-4B/8B and initializes action-head layers from FLUX image-stream weights according to its double-stream and single-stream design.
- Training settings: All models train on 8 NVIDIA H20 GPUs with bf16 precision and AdamW; LIBERO and RoboTwin predict 16-frame-ahead observations with 16-step action chunks.LIBERO uses 10 epochs after resizing concatenated camera views to 224 × 448, while RoboTwin uses 5 epochs with a final 288 × 256 input.
6 Efficiency Optimization
The section evaluates inference-latency optimizations for ImageWAM, including prefix-only attention training, image-denoising-free inference, torch.compile, and static CUDA graphs. It reports latency with three action denoising steps and defines speedup relative to FastWAM using one video denoising step.
- Inference optimization: The evaluation combines prefix-only attention training with image-denoising-free inference to optimize ImageWAM’s inference latency.These strategies are described as similar to those adopted in FastWAM.
- Inference optimization: The study additionally explores torch.compile and static CUDA graphs for model optimization.
- Latency evaluation: 3 action denoising steps are used for all models in the reported inference-latency results.
- Latency evaluation: Speedup is computed relative to FastWAM with one video denoising step.
7 Real-World Experiments Detail
The real-world evaluation covers four challenging manipulation tasks spanning stacking, folding, drawer storage, and hanging a mug. Models use roughly 100 demonstrations per task and are evaluated on 50 trials per task using execution success rate.
- Task Settings: The evaluation includes Stack Three Bowls (T1) and Fold Towel (T2), testing bowl stacking and fabric-towel folding.T1 stacks three green nested bowls; T2 folds a fabric towel.
- Task Settings: Open Drawer & Store Marker (T3) requires opening a drawer, storing a marker, and closing the drawer, while Hang Cup On Rack (T4) hangs a mug on a designated peg.T4 uses a wooden stand as the hanging rack.
- Evaluation Protocol: Execution success rate is the primary performance metric.This metric is used to report real-world task performance.
8 RoboTwin Evaluation Results
This section presents ImageWAM’s per-task RoboTwin evaluation results, reported under both clean and randomized settings.
- Per-task success rates are reported for RoboTwin evaluation in Table 12.
- Table 12 evaluates RoboTwin performance under clean conditions.
- Table 12 also evaluates RoboTwin performance under randomized conditions.