MV-WAM: Manifold-Aware World Action Model with Value Augmentation

2026-06-19

AI Summary

This work addresses the poor policy generalization of embodied agents in out-of-distribution environments, primarily caused by the heterogeneity between visual and action modality manifolds. To tackle this challenge, the authors propose MV-WAM, an end-to-end framework that achieves hierarchical alignment among actions, predicted frames, and value estimates through cross-modal causal masking. The approach further introduces two key innovations: manifold-aware optimization and a progress-value regulation mechanism, which jointly enable autonomous detection of execution deviations and facilitate rollback-based recovery. Evaluated in the RoboTwin simulation without random-action supervision, MV-WAM attains an average success rate of 55.7%, surpassing the strongest baseline by 29.3%. On a real dual-arm robot across four tasks, it achieves an impressive average success rate of 77.5%, substantially narrowing the performance gap between in-distribution and out-of-distribution settings.

Abstract

Achieving robust and generalizable manipulation across diverse environments remains a fundamental challenge in embodied robotics. Recent world action models achieve strong in-domain performance, yet their gains do not extend proportionally to out-of-distribution scenarios. We attribute this to a structural mismatch between visual and action modalities, whose intrinsically heterogeneous manifolds cause joint optimization to disproportionately degrade action robustness under distribution shift. To address this, we propose MV-WAM, a novel end-to-end framework that jointly models visual prediction, action generation, and value estimation designed to effectively leverage video priors during both training and inference for enhanced action generalization. Key to this unification is a cross-modality causal mask that hierarchically grounds actions in predicted video frames and value function tokens in both modalities. To further narrow the generalization gap, MV-WAM adopts a manifold-aware optimization scheme that explicitly accounts for the structural heterogeneity across modalities. Finally, MV-WAM introduces a progress-value regulation mechanism that estimates task completion and detects misalignment between predicted frames and generated actions, enabling the policy to autonomously identify execution deviations and recover through value-guided rollback. On the RoboTwin simulation, MV-WAM achieves a 55.7% mean success rate on random scenarios without any randomized action supervision, outperforming the strongest baseline by 29.3%. MV-WAM achieves a 77.5% mean success rate across four real-world tasks of varying difficulty on a dual-arm robot. Our results demonstrate that manifold-aware cross-modal alignment is essential for robust policy generalization, offering a path toward deployable robotic manipulation.

Problem

• embodied robotics
• action generalization
• distribution shift
• modality heterogeneity
• manifold mismatch

Innovation

• manifold-aware optimization
• cross-modality causal mask
• value-guided rollback
• action generalization
• world action model