DreamWAM: Beyond RGB Future Prediction for World Action Models

2026-08-05

AI Summary

This work addresses the susceptibility of existing RGB-based action models to irrelevant factors such as texture and illumination, which hinders their ability to focus on action-relevant states. To overcome this limitation, the authors propose a structured world modeling approach that jointly represents appearance, motion, geometry, and semantics. During training, a shared attention mechanism guides the action branch to learn relevant features, while at inference time, the model operates solely on RGB inputs. Built upon a shared VideoDiT–ActionDiT architecture, the method integrates an RGB-motion joint latent space denoising process and employs lightweight gated residual branches to encode geometric and semantic information. Evaluated on the LIBERO benchmark, the approach achieves success rates of 98.40% and 98.90% in no-backtrack and joint-inference settings, respectively; on the more challenging LIBERO-Plus, it attains 63.44% and 75.47%, and demonstrates a real-world robot execution success rate of 74.4%, substantially outperforming current baselines.

Abstract

World Action Models (WAMs) learn action-relevant representations by predicting how the observed world will evolve. Most existing WAMs define this future in RGB space, where task-relevant state transitions are entangled with nuisance variations in texture, illumination, background, and viewpoint. We argue that WAMs should explicitly predict action-relevant future state rather than relying on RGB prediction alone. We introduce DreamWAM, which reformulates future prediction as structured world modeling beyond RGB, representing future states through complementary views of appearance, motion, geometry, and semantics. During training, DreamWAM combines joint latent denoising of RGB and motion with lightweight gated residual branches for geometry and semantics. Shared attention between VideoDiT and ActionDiT allows the action branch to learn from these future-state predictions, while all beyond-RGB supervision branches are disabled at inference and deployment remains RGB-only. Across both no-rollout and joint video-action inference, DreamWAM consistently improves the matched RGB-only baselines on LIBERO, from 97.30\% to 98.40\% and from 98.00\% to 98.90\%, respectively. The gains become larger under unseen LIBERO-Plus perturbations, from 51.36\% to 63.44\% and from 69.16\% to 75.47\%. The same robustness extends to real-world manipulation, where DreamWAM attains an average success rate of 74.4\% across unseen changes in lighting, background, and object layout, compared with 55.6\% for Fast-WAM-Joint. These results show that robust world-action learning depends not only on predicting the future, but on representing it in a form that matters for action. The code and models are publicly released at https://github.com/hustvl/DreamWAM.

Problem

• World Action Models
• RGB prediction
• action-relevant representation
• nuisance variations
• future state prediction

Innovation

• World Action Models
• structured world modeling
• beyond-RGB prediction
• latent denoising
• action-relevant representation