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

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.

한 줄 요약

ImageWAM은 비디오 생성 대신 이미지 편집 모델을 활용하여 로봇 행동 예측 성능을 향상시키며, 추론 속도와 계산 비용을 1/4, 1/6 수준으로 줄인다.

핵심 기여도

ImageWAM은 사전 학습된 이미지 편집 모델을 활용하여 로봇 행동 예측을 수행하는 새로운 WAM 프레임워크를 제안한다.
기존 비디오 기반 WAM 대비 FLOPs를 1/6, 추론 지연을 1/4로 줄이며, LIBERO-Plus에서 83.1%의 성공률을 달성한다.
편집 캐시를 활용한 액션 전문가 조건화를 통해, 편집 모델의 내부 표현을 액션 예측에 직접 활용한다.
FLUX.2 4B, Ovis-U1, OmniGen2 등 다양한 편집 모델을 사용 가능하며, 모델 확장 시 성능 개선이 관찰된다.

핵심 아이디어

ImageWAM은 로봇 행동 예측을 위해 이미지 편집 모델을 활용하는 새로운 접근법을 제안한다. 기존 WAM이 비디오 생성을 통해 장기 미래를 예측하는 반면, ImageWAM은 단일 편집 프레임을 기반으로 현재-목표 간 시각적 변화에 집중한다. 이는 액션과 관련된 시각적 변화만을 추출하여, 불필요한 시간적 세부사항을 배제할 수 있다. 편집 모델은 언어 조건에 따라 이미지를 변형하도록 학습되므로, "어떤 변화가 필요한지"를 명확히 이해하는 데 유리하다. ImageWAM은 편집 모델의 내부 표현(편집 캐시)을 액션 전문가에 조건화하여, 명시적 편집 이미지 생성 없이도 액션 예측을 수행한다. 이는 추론 과정에서 비디오 토큰 생성을 생략함으로써 계산 효율성을 높인다.

기술적 접근법

**모델 구조**: ImageWAM은 사전 학습된 이미지 편집 모델(예: FLUX.2 4B, Ovis-U1, OmniGen2)을 사용하며, 편집 모델의 내부 표현(KV 캐시)을 액션 전문가에 조건화한다.
**추론 과정**: 편집 모델은 편집된 이미지를 디코딩하지 않고, 편집 캐시만 생성하여 액션 예측에 활용한다.
**데이터셋**: LIBERO, LIBERO-Plus, RoboTwin 2.0, 실제 세계 조작 작업(4개 태스크)에서 실험을 수행한다.
**하이퍼파라미터**: 모델은 30,000 스텝 동안 훈련되며, 추가 정책 사전 학습 없이 훈련한다.
**모델 아키텍처**: MoT(Multimodal Transformer) 구조를 사용하여 다중 모달 데이터를 결합한다.

주요 결과

**LIBERO-Plus**: FLUX.2 4B 기반 ImageWAM은 평균 성공률 83.1%를 달성하며, FastWAM 및 대부분의 VLA 베이스라인을 초과한다.
**RoboTwin 2.0 (랜덤 환경)**: ImageWAM은 93.56%의 성공률을 기록하며, 모든 VLA 베이스라인을 상회한다.
**성능 개선**: FLOPs는 63.65에서 9.7로, 추론 지연은 1081ms에서 263ms로 각각 1/6, 1/4 수준으로 감소한다.
**모델 확장 효과**: FLUX.2 4B에서 더 큰 FLUX.2 모델로 교체 시, LIBERO-Plus 성공률이 83.1%에서 85.21%로 증가한다.

의의 및 한계

ImageWAM은 비디오 생성 모델의 계산 비용과 불확실성을 줄이며, 이미지 편집 모델의 강력한 언어-시각 조건화 능력을 활용한 새로운 WAM 접근법을 제시한다. 특히, 편집 캐시를 활용한 액션 전문가 조건화는 추론 효율성과 정확도를 동시에 향상시키는 핵심 기술이다. 그러나, 편집 모델의 크기 증가가 모든 퍼트urbation 유형에서 동일한 성능 개선을 보장하지는 않으며, 특정 조건(예: 카메라, 조명, 노이즈)에서는 효과가 제한적일 수 있다. 또한, 편집 모델이 생성하는 변화가 실제 물리적 상호작용을 완벽히 반영하지 못할 가능성도 존재한다.

실용적 활용

ImageWAM은 로봇이 장기적 시각적 변화를 예측하는 대신, 현재-목표 간 변화를 중심으로 행동을 결정해야 하는 조작 작업에 적합하다. 특히, 실시간 제어가 필요한 산업 로봇, 가정용 로봇, 의료 로봇 등에서 계산 효율성과 정확도를 동시에 요구하는 상황에 활용 가능하다. 또한, 다양한 편집 모델과 호환되므로, 이미지 생성 기술이 발전함에 따라 ImageWAM의 성능도 지속적으로 향상될 수 있다.