WAM-Flow: Wasserstein Action Manifolds for Vision-Language-Action Policy Alignment and Optimal-Transport Flow-Matching Decoding

WAM-Flow technical report / preprint. Vision-language-action (VLA) policies commonly decode actions either through discrete tokenization or through continuous Euclidean regression. Both choices can be mismatched to robot control semantics: tokenization introduces irreversible quantization, while Euclidean losses ignore task-dependent temporal, cross-joint, and embodiment structure [Brohan et al., 2023; Kalashnikov et al., 2023; Team et al., 2024; Black et al., 2024]. This paper develops a geometric formulation of action decoding and evaluates it through theory-linked simulations. We define a Wasserstein Action Manifold for action chunks as a task-conditioned Riemannian geometry on trajectory space. In the default construction, the manifold is \( R^T d_a\) equipped with a context-dependent symmetric positive definite metric \(G_ \), where \( =(x, ,s)\) denotes visual, language, and optional proprioceptive context. This metric induces a task-dependent quadratic transport cost and corresponding Wasserstein discrepancy between action distributions [Villani, 2009; Peyre and Cuturi, 2019]. Building on conditional flow matching and optimal transport, we derive a Wasserstein-regularized flow-matching objective for VLA action decoding [Lipman et al., 2023; Tong et al., 2024]. Under explicit assumptions, we state: (i) a coupling-based upper bound showing that discrete decoders inherit quantization distortion controlled by the Wasserstein distance between the true continuous action law and its quantized reconstruction; (ii) a sufficient optimization statement showing that Wasserstein regularization directly penalizes this distortion surrogate; (iii) Lipschitz rollout-sensitivity bounds for the learned flow under observation perturbations; and (iv) a deployment-shift bound in which decoder error depends on Wasserstein distance between training and deployment distributions. Existing OSF archival DOI: 10.17605/OSF.IO/2J765; Existing OSF archival page: https://osf.io/2j765/. Files include the technical report PDF and the LaTeX source tarball when available.