EmbodiedVAE: Disentangled Video VAE for Efficient and Controllable Embodied Manipulation¶

Conference: ECCV 2026
Paper: ECCV 2026
Code: https://github.com/Mutual-Luo/EmbodiedVAE
Area: Image Generation / Robotics & Embodied AI
Keywords: Video VAE / Embodied Manipulation / Action Control / Disentangled Representation / Optimal Transport

TL;DR¶

Addressing the loss of robotic action dynamics caused by generic video VAE compression in embodied manipulation world models, EmbodiedVAE introduces a dual-encoder single-decoder architecture with asymmetric spatio-temporal compression and optimal-transport consistency to disentangle arm motion from background context, achieving a 0.39% compression rate while improving downstream action-conditioned generation by over 2dB PSNR.

Background & Motivation¶

Latent diffusion models (LDMs) have recently shown remarkable promise in embodied learning, particularly for constructing expressive embodied robotic manipulation world models that simulate future robot-environment interactions at minimal trial-and-error cost. However, the latent representations underpinning current world models face an inherent dilemma: relying on 2D image VAEs with frame-by-frame encoding preserves temporal continuity but fails to exploit temporal redundancy across frames, yielding bulky latent representations that severely impede downstream diffusion training efficiency and scalability; conversely, adopting prevailing video VAEs designed for open-domain natural scenes applies uniform spatio-temporal compression without explicit modeling of robotic actions, causing delicate manipulator motion dynamics to be lost in the compressed latent space.

This deficiency is especially critical in action-controlled manipulation video prediction, where the system must faithfully forecast arm pose trajectories and fine-grained object interactions conditioned on future action signals. Existing video VAEs (e.g., Wan-VAE, Cog-VAE) frequently generate visually smooth yet physically inconsistent motions in downstream simulators, exhibiting artificial limb warping or control numbness. This failure stems from the spatial-temporal heterogeneity inherent in embodied setups: the robotic manipulator occupies a localized fraction of the visual scene while governing high-frequency dynamics, whereas the surrounding workspace background is broad yet largely stationary over time. Standard video VAEs downsample the entire video uniformly, inevitably allowing the massive static background to drown out subtle manipulator action features.

This paper tackles the issue by breaking away from uniform scene downsampling and leveraging the dynamic heterogeneity between foreground action and background context. Core idea: design a dual-encoder single-decoder video VAE with asymmetric spatio-temporal compression that automatically disentangles robotic arm motion from environmental context without inference-time masks, regularized by an optimal-transport-based motion consistency loss to enforce temporal coherence on the robotic latent for efficient and controllable manipulation.

Method¶

Overall Architecture¶

EmbodiedVAE consists of a foreground robotic arm encoder, a background environmental context encoder, and a unified reconstruction decoder featuring dual-head upsampling and a shared tail module. During inference, the system operates completely without external segmentation masks or sensor prompts: raw manipulation video clips are automatically mapped into two complementary, compact latent representations—one explicitly capturing fine-grained manipulator dynamics, and the other modeling global static scene context. The unified decoder then aligns and merges both latent streams to reconstruct the full high-fidelity video.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Manipulation Video<br/>X in R^(T×H×W×C)"] --> B["Dual-Encoder Asymmetric Spatio-Temporal Downsampling<br/>AlphaBlender 2D+1D Hybrid Blocks"]
    B --> C["Foreground Robotic Arm Latent Z_arm<br/>4×(16×16) Heavy Spatial Compression"]
    B --> D["Background Environmental Latent Z_env<br/>8×(8×8) Heavy Temporal Compression"]
    C --> E["OT Motion Consistency Constraint<br/>Sinkhorn Algorithm Minimizing Transport Cost"]
    C --> F["Unified Reconstruction Decoder<br/>Dual-Head Upsampling → Concatenation → Shared Tail"]
    D --> F
    F --> G["Reconstructed Video<br/>X_hat"]

Key Designs¶

1. Dual-Encoder Asymmetric Spatio-Temporal Compression: Custom Downsampling for Heterogeneous Dynamics

Generic video VAEs fail in robotic tasks because they apply uniform spatial and temporal downsampling factors across locally dynamic objects and globally static backgrounds. Recognizing that the manipulator covers a compact spatial area with rapid transitions while the workspace background exhibits immense temporal redundancy, EmbodiedVAE constructs two tailored encoding paths. To balance representational power with compute efficiency, the encoders avoid heavy 3D convolutions, applying causal 3D convolutions only at the input and output layers and stacking hybrid blocks of 2D spatial and causal 1D temporal downsampling in between, blended via the AlphaBlender operator: $$ \mathbf{x}_{\text{down}} = \alpha \cdot \mathbf{x}_1 + (1 - \alpha) \cdot \mathbf{x}_2 $$ where $\mathbf{x}_1$ and $\mathbf{x}_2$ denote features from the interpolation and convolution downsampling branches respectively, with $\alpha = \text{Sigmoid}(0.2)$ fixed. The robotic arm encoder stacks 4 spatial and 2 temporal downsampling blocks to achieve a $4\times(16\times16)$ (temporal $\times$ spatial) compression, restricting temporal subsampling to 4$\times$ to prevent motion aliasing while aggressively compressing spatial dimensions. In contrast, the environment context encoder stacks 3 spatial and 3 temporal downsampling blocks for $8\times(8\times8)$ compression, purging background temporal redundancy. Both latent spaces use 4 channels, achieving an overall compression rate of 0.39% while cleanly preserving fine-grained action semantics.

2. Optimal-Transport-Based Motion Consistency: Explicit Inter-Frame Dynamics Alignment

Compressing the temporal dimension can cause subtle inter-frame arm displacements to be smoothed out or lost, damaging latent temporal coherence. Standard pairwise feature metrics capture only local patch affinities rather than global structural motion. Grounded in the principle that stable motion dynamics exhibit visual correspondence invariance across consecutive frames, EmbodiedVAE models latent slices $\mathbf{Z}_p$ and $\mathbf{Z}_q$ at timestamps $t_p, t_q$ as discrete probability distributions and poses an entropy-regularized optimal transport problem with negative cosine similarity as transport cost $c(i,j) = -\langle \mathbf{Z}_p[i], \mathbf{Z}_q[j] \rangle / \sqrt{d}$: $$ \mathbf{T}^\star = \arg\min_{\mathbf{T}} \sum_{i,j} \mathbf{T}{ij} c(i,j) + \epsilon \mathcal{H}(\mathbf{T}), \quad \text{s.t.} \quad \mathbf{T}\mathbf{1} = \mu_p, \ \mathbf{T}^\top\mathbf{1} = \mu_q $$ The differentiable transport plan $\mathbf{T}_{p\rightarrow q}$ is solved via the Sinkhorn-Knopp algorithm, yielding the motion consistency loss: $$ \mathcal{L}}} = \frac{1}{N} \sum_{p=1}^{N} \sum_{q=p}^{p+\Delta} |\mathbf{Zq - \mathbf{T}|_F $$ This loss aligns latent feature transitions under the optimal transport mapping, while the permutation penalty } \mathbf{Z}_p|_2 + \gamma |\mathbf{T}\mathbf{T}^\top - \mathbf{I$\|\mathbf{T}\mathbf{T}^\top - \mathbf{I}\|$ encourages one-to-one correspondences. Mathematically, this enforces temporal smoothness and prevents abrupt latent drifts across frames.

3. Two-Stage Mask-Guided Disentangled Training: Zero Inference-Time Mask Overhead

Direct end-to-end training of dual encoders tends to cause feature leakage and entangled representations where both encoders capture redundant scene information. EmbodiedVAE resolves this via a two-stage training strategy: in the first stage, auxiliary binary arm masks $\mathbf{M}$ are generated offline using SAM2 and Grounding DINO with the prompt "robotic arm". The arm encoder is trained to reconstruct only foreground regions $\mathcal{L}_{\text{rec}} = \|(\hat{\mathbf{X}} - \mathbf{X}) \odot \mathbf{M}\|$ augmented with $\mathcal{L}_{\text{motion}}$, while the environment encoder reconstructs the background $\mathcal{L}_{\text{rec}} = \|(\hat{\mathbf{X}} - \mathbf{X}) \odot (1-\mathbf{M})\|$. In the second stage, the separate branch decoders are discarded, the two encoders are frozen, and the unified decoder is trained to reconstruct the complete video $\mathcal{L}_{\text{rec}} = \|\hat{\mathbf{X}} - \mathbf{X}\|$. Consequently, the model operates completely mask-free during inference, relying on internal inductive biases learned during training.

Loss & Training¶

In the first training stage, the arm encoder objective is: $$ \mathcal{L}{\text{arm}} = \mathcal{L}}}(\mathbf{M}) + \alpha \mathcal{L{\text{motion}} + \beta \mathcal{L}}} + \lambda \mathcal{L{\text{aux}} $$ and the environment encoder objective is $\mathcal{L}_{\text{env}} = \mathcal{L}_{\text{rec}}(1-\mathbf{M}) + \beta \mathcal{L}_{\text{KL}} + \lambda \mathcal{L}_{\text{aux}}$, where $\mathcal{L}_{\text{aux}}$ combines feature-level perceptual loss and adversarial loss. In the second stage, the unified decoder objective is: $$ \mathcal{L}|}} = |\hat{\mathbf{X}} - \mathbf{X1 + \beta \mathcal{L} $$ When integrated into downstream world models like IRASim-L (461M), the latent tokens of }} + \lambda \mathcal{L}_{\text{aux}$\mathbf{Z}_{\text{arm}}$ and $\mathbf{Z}_{\text{env}}$ are processed through shared DiT blocks, with cross-attention inserted every $k=2$ blocks to enable adaptive physical interaction between arm dynamics and environment illumination/contact.

Key Experimental Results¶

Main Results¶

The authors evaluated EmbodiedVAE on video reconstruction using Agibot-2025 and Bridge datasets (Table 1), and on action-conditioned video generation using RT-1 and Bridge datasets (Table 2).

Table 1: Video VAE reconstruction quality comparison (Original Paper Table 1)

Method	Com. Rate↓	Agibot-2025 PSNR↑	Agibot-2025 SSIM↑	Agibot-2025 FVD↓	Bridge PSNR↑	Bridge SSIM↑	Bridge FVD↓
OpenSoraPlan	1.04%	30.7553	0.9257	290.2550	30.3755	0.8948	340.5284
Cosmos	2.08%	28.2329	0.8778	889.7757	27.8812	0.8386	715.6345
iVideoGPT	1.50%	29.2503	0.9096	760.6489	28.3325	0.8738	522.1818
MAGVIT-v2	0.65%	27.1014	0.7197	849.5063	26.8635	0.6986	549.9939
VidTwin	0.20%	30.8263	0.9202	443.1976	29.7386	0.8531	585.8065
EMU-3	0.53%	28.9081	0.8926	599.8921	27.5814	0.8317	736.5338
CV-VAE	0.53%	30.2767	0.9128	392.6846	29.8419	0.8789	418.8394
CMD	6.85%	31.4030	0.8995	439.9010	31.3764	0.8996	358.6915
EmbodiedVAE (Ours)	0.39%	31.6745	0.9345	368.5511	30.6957	0.9017	309.7604

Table 2: Downstream action-controlled manipulation video generation comparison (Original Paper Table 2, backbone fixed to IRASim-L)

Method	Com. Rate↓	RT-1 PSNR↑	RT-1 SSIM↑	RT-1 FVD↓	Bridge PSNR↑	Bridge SSIM↑	Bridge FVD↓
SDXL (2D Image VAE)	8.32%	23.1988	0.8074	761.7837	22.3624	0.7915	455.7345
OpenSoraPlan	1.04%	21.6494	0.7432	958.0784	21.0405	0.7577	785.5559
EMU-3	0.53%	21.6749	0.7286	1049.2935	20.8573	0.7053	986.3582
CV-VAE	0.53%	20.3296	0.6859	2529.3193	19.3068	0.6568	1829.7292
CMD	6.85%	20.9454	0.7671	2644.7208	20.1799	0.7419	1464.6843
VidTwin	0.20%	21.1037	0.6697	1824.3383	20.4466	0.6889	1690.1818
Cog-VAE	2.08%	21.2231	0.7907	1135.6917	21.6951	0.7807	1302.4434
Wan-VAE	2.45%	22.7334	0.8138	802.2071	22.1690	0.8211	548.9621
EmbodiedVAE (Ours)	0.39%	24.7082	0.8263	716.3273	24.2484	0.8366	631.6450

Ablation Study¶

Table 3: Ablation studies on reconstruction and downstream action-controlled manipulation (Original Paper Table 3)

Task	Config	Agibot-2025 PSNR↑	Agibot-2025 LPIPS↓	Bridge PSNR↑	Bridge LPIPS↓
Recon.	EmbodiedVAE (full model)	31.6745	0.0723	30.6957	0.1060
Recon.	w/o Mask (no two-stage mask guide)	31.1812	0.0728	29.6680	0.1211
Recon.	w/o $\mathcal{L}_{\text{motion}}$ (no OT consistency loss)	31.0213	0.0835	29.9769	0.1087
Task	Config	RT-1 PSNR↑	RT-1 LPIPS↓	Bridge PSNR↑	Bridge LPIPS↓
Manip.	EmbodiedVAE (full model)	24.7082	0.1707	24.2484	0.1409
Manip.	w/o Mask (no two-stage mask guide)	23.8044	0.1904	23.7306	0.1705
Manip.	w/o $\mathcal{L}_{\text{motion}}$ (no OT consistency loss)	23.8187	0.1863	23.9726	0.1643

Key Findings¶

Superior action controllability over frame-by-frame 2D VAE: In RT-1 and Bridge downstream action-controlled generation, EmbodiedVAE with only 0.39% compression outperforms uncompressed 2D SDXL (8.32% footprint) by 1.51 dB and 1.88 dB in PSNR, respectively, and outclasses the runner-up video VAE (Wan-VAE) by an average of 2.02 dB PSNR.
Disentanglement and OT consistency are vital for downstream action control: Removing $\mathcal{L}_{\text{motion}}$ drops reconstruction PSNR on Agibot by 0.65 dB, but causes a sharper 0.89 dB drop on downstream RT-1 manipulation generation; removing two-stage mask guidance induces a similar ~0.9 dB drop on RT-1, demonstrating that decoupling arm representations from background noise is crucial for action injection.
Dramatic reduction in downstream world model training overhead: Pre-encoded latent benchmarks show that downstream LDM training with EmbodiedVAE is nearly 4$\times$ faster than with SDXL latents over 1000 steps, dramatically cutting memory footprint and computational cost.

Highlights & Insights¶

Asymmetric spatio-temporal compression matching scene physics: Allocating 4$\times$ temporal / 16$\times$ spatial compression to the manipulator and 8$\times$ temporal / 8$\times$ spatial to the background elegantly resolves the trade-off between action dynamic fidelity and overall latent compactness.
Mask-guided training with zero-mask inference: Supervising disentanglement via SAM2/Grounding DINO during training while discarding auxiliary decoders in stage two enables mask-free deployment without extra segmentation inference latency.
Global optimal transport regularizer for temporal dynamics: Formulating inter-frame latent feature tracking as an entropy-regularized optimal transport problem provides a mathematically sound bijective mapping that guarantees smooth, non-collapsing motion dynamics.

Limitations & Future Work¶

Reliance on relatively stationary camera view: The aggressive 8$\times$ temporal compression on environmental background assumes fixed or gently moving first-person perspectives typical of tabletop manipulation; it may exhibit blurring under high-speed mobile manipulation camera motions.
Single-manipulator semantic prior: The first-stage mask supervision focuses specifically on robotic arms; generalizing to complex multi-agent bimanual setups, highly articulated multi-finger dexterous hands, or deformable object state transitions requires richer segmentation priors.
Future direction: Extending the disentanglement pipeline into a tripartite representation (manipulator, manipulated object, and background context) to explicitly track contact mechanics.

vs. Wan-VAE / Cog-VAE: General-purpose video VAEs optimize for open-domain visual fidelity and smooth spatio-temporal pooling, diluting localized action dynamics; EmbodiedVAE achieves superior action-conditioned generation quality (+2.02 dB PSNR) at a fraction of the latent size (0.39% vs. 2.08%~2.45%).
vs. VidTwin / CMD: VidTwin and CMD explore structural-dynamic decoupling in natural scenes via unsupervised heuristics, which often miss small-scale robotic manipulators against cluttered backgrounds; EmbodiedVAE leverages domain-specific asymmetric compression and OT consistency for reliable action control.

Rating¶

Novelty: ⭐⭐⭐⭐☆ (Strong domain-tailored insight in designing asymmetric spatio-temporal compression and OT-based motion consistency for embodied manipulation)
Experimental Thoroughness: ⭐⭐⭐⭐⭐ (Trained on 1M+ manipulation video clips; evaluated on two reconstruction and two downstream world model generation benchmarks with extensive ablations)
Writing Quality: ⭐⭐⭐⭐⭐ (Well-structured paper with rigorous mathematical derivations and comprehensive visual evidence)
Value: ⭐⭐⭐⭐⭐ (Delivers an essential high-compression, action-controllable video latent tokenizer for next-generation embodied world models)