EmbodiedVAE: Disentangled Video VAE for Efficient and Controllable Embodied Manipulation

摘要

Latent diffusion models (LDMs) have recently significantlyadvanced embodied learning in constructing powerful embodied manipu-lation world models. However, despite the remarkable performance, exist-ing LDMs predominantly rely on Variational Autoencoders (VAEs) opti-mized for natural scenes while failing to account for the unique character-istics of embodied manipulation scenarios, yielding latent representationsthat are neither compact nor controllable, thereby hindering efficienttraining of LDMs and precise robotic control. To solve this problem, wepresent EmbodiedVAE, a novel video VAE that provides compact yetcontrollable latent representations tailored for the robotic manipulationworld models. Specifically, EmbodiedVAE adopts a dual-encoder, single-decoder architecture with an asymmetric spatio-temporal compressionmodule, which automatically disentangles the robot arm’s motion frombackground environment, resulting in overall compactness while provid-ing explicit embodied latent to support fine-grained action control. Tofurther preserve the temporal consistency of learned robotic motion la-tent, we introduce an optimal-transport-based consistency module thatexplicitly enforces motion fidelity and inter-frame coherence. Extensiveexperiments demonstrate that our proposed EmbodiedVAE achieves su-perior reconstruction quality with high compression rate, while enablingmore precise action control in robotic manipulation scenarios with anaverage of 2dB PSNR improvement over state-of-the-art video VAEs.Code is available at: https://github.com/Mutual-Luo/EmbodiedVAE