Latent embeddings improve action generation for robot manipulation tasks
Latent evolving World Action Model
Computer Vision and Pattern RecognitionRobotics
Summary
Generating robot actions requires understanding how the environment changes over time. This paper shows that using special visual representations called JEPA embeddings helps make better action predictions than older methods. The authors created LeWAM, a system that predicts future robot actions using these embeddings without needing expensive video models. They also developed a way to improve action quality by learning preferences directly from examples, boosting performance in simulated and real-world tasks.
What this means in practice
- •For robotics engineers: Develop robot controllers that generate effective actions by predicting environment changes using latent visual embeddings instead of large video models.
- •For automation system developers: Implement offline preference refinement to improve action quality from demonstration data without extra environment interaction or human oversight.
Authors
Xueji Fang, Boqiang Duan, Hua Wu, Jingdong Wang, Guo-Jun Qi
Abstract
World Action Models (WAMs) jointly model action generation and environment dynamics and are mostly built on pretrained Video Diffusion Models (VDMs). In VDM-based WAMs, observations are first encoded by a VAE, and the resulting compressed latents are then processed by large video diffusion backbones to extract effective features for action generation. However, this paradigm ties WAM performance and training cost to large-scale video generation pretraining, limiting WAM efficiency and scalability. In this paper, we theoretically and empirically investigate how visual representations affect action generation in WAMs. Our results show that predictive embeddings from Joint-Embedding Predictive Architecture (JEPA) encoders better support action generation than compressed VAE latents, with I-JEPA performing best in our encoder comparison. Based on these findings, we propose LeWAM, which conditions action generation on JEPA embeddings and models environment evolution by predicting future embeddings in the same space, without relying on a video diffusion backbone. We further find that imitation learning matches demonstrated actions but does not distinguish better actions from worse ones, even though small action deviations can greatly affect task success. To address this limitation without additional environment interaction or the human oversight required for resets and safety, we introduce Demonstration-Guided DPO (DemoDPO), an offline preference refinement stage that derives preference supervision directly from demonstrations.With only 0.4B trainable parameters, LeWAM achieves an average success rate of 92.28\% on RoboTwin 2.0, comparable to that of state-of-the-art VLAs and WAMs, and maintains practical effectiveness on real-world manipulation tasks.