ImageWAM: Efficient World Action Modeling

• ImageWAM is a world action modeling framework that replaces full video generation with endpoint visual edits, using image-editing diffusion models to predict task-driven actions.
• It extracts compact KV-caches from transformer layers to focus solely on action-relevant scene changes, significantly reducing computational load.
• Empirical evaluations on simulated and real-world tasks demonstrate substantial efficiency gains and improved success rates compared to traditional video-based approaches.

ImageWAM is a world action modeling (WAM) framework for language-conditioned robot control that replaces standard video generation with pretrained image-editing diffusion models, leveraging their ability to represent and ground task-driven visual transformations. Rather than generating full video rollouts, ImageWAM extracts compact representations from the internal states of image-editing transformers and conditions a diffusion-based flow-matching policy solely on these features to predict actions. This approach yields substantial efficiency gains and improves action relevance by restricting world modeling to action-centric visual changes (Zhang et al., 17 Jun 2026).

1. Framework Overview and Motivation

Traditional WAMs in robot learning rely on future video generation as the backbone for perceiving, modeling, and forecasting the effects of actions in the world. However, video-based approaches incur high computational cost due to dense tokenization across time, allocate model capacity toward background and appearance modeling irrelevant to planning, and suffer from compounding temporal errors in long-horizon prediction. ImageWAM addresses the central question: must a WAM generate future videos, or can it model only the "endpoint" transformation—i.e., the visual edit from present to immediate goal state?

ImageWAM employs an image-editing denoising model (typically a text-conditioned diffusion model), which, given the current observation $o_t$ and an instruction $\ell$, predicts the latent of the single future "target" frame $o_{t+H+1}$. Instead of actually synthesizing the target image at inference, ImageWAM intercepts key/value ("KV") caches from the model's transformer blocks at a chosen denoising step, producing a compact representation $C_{\rm edit}$ focused on scene changes pertinent to the task. This cache informs a flow-matching action expert responsible for generating the action chunk $a_{t:t+H}$.

2. Architectural Components and Learning Objectives

ImageWAM consists of three essential modules:

1. Image-Editing Denoiser: Takes $(o_t, \ell)$ and targets the VAE latent $z_{t+H+1}$ of the goal frame. Training uses the score-matching loss:

$$L_{\rm img} = \mathbb{E}_{z_{t+H+1}, \epsilon, r} \left\| \epsilon - \epsilon_\theta\left(z_r, r; o_t, \ell\right) \right\|^2$$

where $z_r = \sqrt{1-r}z_{t+H+1} + \sqrt{r}\,\epsilon$.

2. KV-Cache Extraction: During training, a random denoising step $T$ is sampled, and all transformer layer caches after cross-attention with $\ell$0 and image tokens are extracted:

$\ell$1

This cache is a compact, task-grounded summary encoding intended transformation.

3. Flow-Matching Action Expert: Given $\ell$2 and context, the expert predicts action denoising velocities:

$\ell$3

where $\ell$4.

The total objective is $\ell$5.

At inference, only a single forward denoising pass at a fixed $\ell$6 is performed to obtain $\ell$7; no image decoding or multi-frame rollout occurs. The action expert samples the action chunk using efficient flow-matching denoising steps.

3. Computational Efficiency and Capacity Allocation

ImageWAM's parameter and compute efficiency directly derive from its focus on endpoint visual edits. The method reduces inference latency and FLOPs by avoiding multi-frame future generation:

Method	Latency (ms)	TFLOPs
Video-WAM (full rollout)	1081	63.65
Fast-WAM (1-step)	302	13.21
ImageWAM	263	9.72

Crucially, video-based WAMs must allocate model capacity to background, appearance, and temporal continuity (e.g., camera motion and unrelated dynamics), whereas ImageWAM's edit prior is inherently localized to action-relevant regions and grounded in the instruction. This targeted allocation mitigates overfitting on irrelevant dynamics and reduces the risk of compounding errors in long-horizon tasks.

4. Experimental Evaluation and Benchmarking

ImageWAM's empirical performance has been validated on:

• Simulated environments: RoboTwin 2.0 (bimanual, 50+ tasks, clean/randomized), LIBERO (Spatial, Object, Goal, Long suites), LIBERO-Plus (perturbed visuals).
• Real-world robotics: Dobot XTrainer (dual-arm) on long-horizon and fine-motor tasks (e.g., stacking bowls, folding towels, opening drawers).

Key benchmarking results:

Environment	Method	Success Rate (%)
RoboTwin (Clean)	VLA Baselines	76–88
	Fast-WAM	91.9
	ImageWAM	93.2
LIBERO-Plus	Fast-WAM	51.5
	ImageWAM	83.1 (Flux2-4B)
Real-World (Avg)	π₀.₅	72.3
	Fast-WAM	79.0
	ImageWAM	84.5

Notably, ablations indicate Flux2-9B edit backbones further push LIBERO-Plus averages to 85.2%; decoupled architectures where image editing and action policy are separated outperform unified VLA backbones (e.g., 98.4% vs. 84.4% on LIBERO).

5. Attention, Interpretability, and Failure Modes

Cross-attention analysis for the flow-matching expert reveals that in ImageWAM, attention is highly concentrated on objects being manipulated, target receptacles, or candidate contact zones. Backgrounds are correspondingly de-emphasized, indicating that the underlying cache features robustly localize to intended visual changes dictated by the instruction. In contrast, baseline video-WAM attention is often distributed across irrelevant regions or the entirety of generated frames.

Limitations emerge in tasks requiring multi-step visual planning, large viewpoint shifts, or persistent temporal feedback where a single image edit is insufficient as a world model. In highly dynamic scenes with rich temporal dependencies, classical video generation may remain superior.

6. Broader Implications and Prospects

ImageWAM introduces a paradigm shift for robot world modeling, indicating that modeling the task-driven endpoint transformation is often sufficient for instruction-conditioned action prediction—eliminating much of the cost and modeling complexity associated with video generation. Computational requirements are reduced by factors of 4–6× in practice. Moreover, ImageWAM's decoupled architecture enables integration with a diverse range of image-editing backbones and facilitates future enhancements, including multi-step editing chains, hybrid world models with partial video rollouts, fusion with geometric features (e.g., depth or segmentation cues), and policy fine-tuning.

Adoption of compact, interpretable world-action representations grounded in editing caches promises broader scalability to real-world robotics contexts with diverse visual requirements (Zhang et al., 17 Jun 2026).