ImageWAM: Do World Action Models Really Need Video Generation, or Just Image Editing?

Instead of making a robot imagine a whole video of the future to decide its next move, this new method uses an image editing model to just "imagine" the final result. This makes the robot much faster and more accurate. It does this because it focuses only on the changes that matter for the task.

In the field of robotics, World Action Models (WAMs) have become a dominant paradigm. These models attempt to bridge the gap between seeing the world and acting within it. They use generative models to "imagine" future states. The prevailing logic is that if a model can accurately predict a video of what happens next, it can derive the necessary motor commands. This "reason-before-act" loop relies on video generation to provide a rich, temporal prior of physics and object motion.

However, this reliance on video is increasingly looking like a massive computational tax for very little return. Recent research suggests that predicting every pixel of a multi-frame video sequence is an inefficient way to solve a manipulation task. A new paper, ImageWAM, challenges this status quo. It asks if we can replace dense video generation with something much leaner: image editing.

The overhead of imagining everything

Current video-based WAMs suffer from a fundamental mismatch between the generative objective and the robotic task. To predict a future video, a model must spend significant capacity modeling things that are irrelevant to the robot's next action. Examples include background textures, lighting shifts, or camera jitter. As shown in [Figure 1(a)], video generation forces the model to synthesize a dense stream of spatio-temporal tokens (units of information representing data across space and time). This creates three distinct failure modes.

First, there is the sheer cost of inference (the process of a trained model making predictions). Generating multi-frame sequences is slow. This is a dealbreaker for real-time closed-loop control. Second, the model wastes its limited parameter budget on "appearance details" rather than the critical delta (difference) between the current state and the goal. Finally, there is the risk of "hallucinated" physics. If a model generates a distorted or physically impossible video frame, the downstream action predictor will ingest that error. This leads to catastrophic movements. The authors illustrate this in, where artifacts in the imagined video lead to failed grasp attempts or incorrect object placement.

Repurposing the editing prior

The core insight of ImageWAM is that robot manipulation is essentially an instruction-guided visual transformation. When you tell a robot to "pick up the cup," you aren't asking it to simulate the entire history of the universe. You are asking it to transform the current scene into one where the cup has moved. This is exactly what pretrained image editing models are designed to do.

The architecture follows a specific pipeline :

1. Feature Extraction: The model takes the current observation ($o_t$) and the language instruction ($l$). It passes them through a pretrained image editing backbone (such as OmniGen2, Ovis-U1, or FLUX.2).
2. KV Cache Harvesting: Crucially, the model does not actually decode the edited image at inference time. Instead, it intercepts the intermediate Key-Value (KV) caches produced during the denoising process. These caches represent the model's internal "understanding" of how the image should change to satisfy the instruction.
3. Action Expert Conditioning: These KV caches are treated as a compact world-action context. They are fed into a separate "Action Expert"—a flow-matching transformer (a type of generative model that learns to map noise to data along a straight path)—which uses this context to predict the required action chunk ($a_{t:t+H}$).

By using the KV caches instead of the decoded pixels, the model bypasses the need to render a high-fidelity target frame. This effectively turns a generative model into a highly efficient feature extractor for intent.

Efficiency without the sacrifice

The empirical results suggest that this shift from video to editing is not just a theoretical win. It translates directly to the bottom line of deployment. The authors report that ImageWAM achieves a success rate of 93.56% on the RoboTwin 2.0 randomized benchmark. This outperforms standard VLA (Vision-Language-Action) baselines. It also matches or exceeds competitive WAMs.

For an engineer concerned with the hardware footprint, the efficiency gains are the headline. The models were trained on 8 NVIDIA H20 GPUs. The paper reports that ImageWAM reduces FLOPs (floating-point operations, a measure of computational work) to 1/6 and latency (the delay before a response) to 1/4 compared to video-based WAMs. Looking at the specific breakdown in [Table 5], a standard video-based WAM might hit latencies north of 1000ms. This is too slow for fluid motion. In contrast, ImageWAM operates in the 263ms range on an A6000 GPU. This allows for much more responsive control.

Furthermore, the model shows unexpected robustness in complex scenarios. In real-world experiments using a dual-arm Xtrainer robot, ImageWAM achieved an 84.5% success rate. This notably outperformed $\pi$0 (55.8%) and $\pi$0.5 (72.3%). The authors note that the gains were most pronounced in tasks involving deformable objects, like folding a towel. This suggests that the "change-centric" nature of editing caches is particularly good at capturing the nuances of non-rigid manipulation.

The limits of scaling and unification

Despite the strong performance, there are edges to this approach. First, the relationship between backbone scale and robustness is not perfectly linear. While moving from FLUX.2 4B to 9B improved success rates in several categories, the improvements were not uniform. The authors found that improvements in camera angle, lighting, and noise robustness were not monotonic [Table 7]. This means that simply adding more parameters won't automatically solve every environmental perturbation.

Second, the paper touches on the tension between "understanding" and "generation." The authors chose to keep the VLM (Vision-Language Model) components frozen. They only train the diffusion branch and action expert. They argue that unified models can suffer from interference. This occurs when the requirements for high-level semantic abstraction clash with the requirements for fine-grained spatial generation [Table 6]. While decoupling works here, it raises questions about future architectures.

Verdict: Ship the cache, skip the video

If you are building a real-time robot policy, the takeaway is clear. Stop trying to force your model to render high-fidelity future videos just to get a sense of intent. The "imagination" required for control is much closer to an image edit than a cinematic rollout.

ImageWAM proves that the intermediate representations (the KV caches) of a generative model are far more valuable for action prediction than the final decoded pixels. The trade-off is incredibly favorable. You get comparable success rates with a 4x reduction in latency and a massive drop in compute requirements.

The methodology is sound. The implementation details—including the use of flow-matching for the action expert and the weight-copy initialization strategy—are documented well. Code is reportedly available; see the paper for the canonical link. If you are currently struggling with the latency tails of a video-based WAM, this is a high-priority architectural pivot.

Figures from the paper