Flowing With Purpose: Latent Action Guided Flow Matching Policies For Robotic Manipulation

Abstract: Flow matching has recently become a new standard for behavior cloning in robotic manipulation. However, state-of-the-art flow matching policies suffer from a systematic structural mismatch: they rely on a globally fixed isotropic source distribution despite the strongly fragmented and heteroscedastic structure of robotic action spaces. This agnostic initialization forces the model to learn highly entangled vector fields, bottlenecking training efficiency and limiting overall policy performance. To address this limitation, we introduce Latent Action Guided Flow Matching (LAFM), a novel framework that replaces the monolithic Gaussian with an adaptive library of learned prior distributions. By grounding these distributions using a latent action model, LAFM maps current observations to discrete motion primitives, selecting a specialized base distribution that provides an informed, structurally aligned initialization for the denoising process. This dynamic adaptivity naturally accommodates heteroscedasticity in human demonstrations and makes transport trajectories shorter and less entangled. Empirically, LAFM substantially outperforms standard flow matching formulations, increasing task success rates by 23.4% in real-world robotic deployments and by 10.4% on the LIBERO-90 benchmark. Furthermore, we demonstrate that LAFM achieves state-of-the-art results, surpassing massively pre-trained vision-language-action models while utilizing significantly smaller architectures.

• The paper introduces latent action guided flow matching (LAFM) that replaces fixed isotropic priors with a learned library of Gaussian distributions to reduce transport cost.
• It employs a multimodal encoder and transformer to align noise sampling with discrete motion primitives, resulting in smoother denoising trajectories.
• Experimental results on Franka Emika Panda and LIBERO benchmarks demonstrate significant success rate improvements over standard flow matching methods.

Latent Action Guided Flow Matching Policies for Robotic Manipulation

Introduction and Motivation

Recent advances in visuomotor policy learning have adopted generative denoising paradigms such as diffusion and flow matching policies for robotic manipulation. Flow matching, in particular, provides efficient inference and stable optimization by learning a continuous vector field that transports samples from a base distribution to the action distribution. However, prevailing flow matching methods initiate all denoising trajectories from a fixed isotropic Gaussian prior, which fails to account for the fragmented and heteroscedastic nature of robotic action spaces. This induces structural bottlenecks, leading to highly entangled vector fields and inefficient transport trajectories, especially in multi-modal settings where similar observations correspond to disjoint action modes.

To address this, the paper proposes Latent Action Guided Flow Matching (LAFM), which replaces the monolithic prior with a library of adaptive learned distributions, grounded using a Latent Action Model (LAM). The LAM transforms current observations into discrete motion primitives, aligning prior initialization with coarse action intent and naturally accommodating the heteroscedasticity of human demonstrations.

Figure 1: LAFM overview—LAM extracts latent actions, policy encoder predicts the next latent action, which conditions the noise sampling and guides the flow matching decoder.

Methodology

Flow Matching Policy Review

The standard flow matching policy optimizes a vector field that transports noise samples from $\mathcal{N}(\mathbf{0},\mathbf{I})$ to the target action distribution, minimizing the squared error between predicted and true denoising vectors. The transport probability path is defined by linear interpolation between action targets and noise, and inference is performed via Euler integration over flow steps.

Adaptive Priors via Latent Actions

LAFM introduces a library of $K$ learned Gaussian priors $$\{\mathcal{N}(\boldsymbol{\mu}_k,\boldsymbol{\Sigma}_k)\}$$, indexed by discrete latent actions predicted from current observations. The categorical latent action index is selected via a pre-trained LAM encoder, and forms an explicit inductive bias for generative policy modeling. The conditional flow matching objective adapts the denoising vector field to the latent-conditioned prior, reducing expected transport distance and vector field entanglement.

A formal proposition establishes that, under mode-aligned priors with reduced covariance, the expected squared transport distance for latent-conditioned priors is strictly less than that for a single isotropic prior, validating the geometric intuition underpinning LAFM.

Architecture and Regularization

The proposed architecture comprises an encoder-decoder structure: image, language, and proprioception inputs are processed with a specialized encoder and transformer; the encoder predicts the latent action, which indexes the prior for noise sampling; the decoder (DiT) operates on the noisy action sequence with cross-attention. All priors are initialized to standard normal for stable training, and their specialization is controlled via KL-divergence regularization, allowing interpolation between single-prior and multi-prior regimes.

Figure 2: Architecture: multimodal encoder feeds transformer, prompt token yields latent action, which selects prior embeddings for noise—cross-attention conditions DiT decoder.

Experimental Evaluation

Real-World Robotic Manipulation

LAFM was evaluated on four manipulation tasks using a Franka Emika Panda robot, with policies compared against FM, ACT, and $\pi_0$ baselines. LAFM achieved an average success rate improvement of 23.4% over FM and outperformed $\pi_0$ (which uses extensive VLA policy pretraining) by 15% in success rate and 11% in completion score. Empirically, LAFM produces smoother, more direct transport trajectories within real-world rollouts.

LIBERO Benchmarks

On the LIBERO-90 suite (90 tasks, 3,959 episodes), LAFM yields a 10.4% absolute success rate improvement over FM and outperforms all other reported baselines, including state-of-the-art vision-language-action models, while requiring fewer parameters and limited pretraining. On the combined LIBERO Spatial, Object, Goal, and Long benchmarks, LAFM consistently outperformed FM and other baselines, exhibiting strong generalization and competitive performance against larger pre-trained counterparts.

Auxiliary Supervision and Ablations

Analysis indicates that latent action prediction supervision improves FM performance, but structured prior initialization in LAFM accounts for the full performance gain, confirming the importance of mode-aligned source distributions. Ablations on LAM codebook size and training data demonstrate that both under- and over-partitioning degrade performance, but LAFM consistently outperforms FM across codebook choices and data sources, highlighting the importance of prior quality and coarse motion diversity.

Figure 3: Impact of LAM codebook size and training dataset on LIBERO-90 success rate—optimal partitioning yields maximal performance, with consistent superiority over standard flow matching.

Figure 4: Success rate on LIBERO-90 as a function of denoising steps—LAFM maintains robust performance across steps, demonstrating inference efficiency.

Qualitative Analysis

Learned Prior Distributions

t-SNE visualization of learned prior means reveals that similar latent actions cluster, with priors for large-range motions exhibiting higher variance. Optical flow overlays confirm that the learned priors correspond to distinct motion primitives, reinforcing the suitability of latent actions for capturing heteroscedasticity in demonstration data.

Figure 5: t-SNE of learned prior means and variances, with optical flow overlays for representative latent actions—priors align with motion primitive structure and variance.

Flow Matching Vector Fields

Visualization of denoising vector fields for FM and LAFM on LIBERO-90 batch data shows that latent-conditioned priors produce shorter, less entangled transport trajectories, materially reducing transport cost and prediction error. PCA projection empirically supports Proposition 1: mode-aligned priors streamline the transport geometry and improve flow matching efficiency.

Figure 6: Comparison of denoising vector fields for FM and LAFM—gray dots are noisy samples, blue lines are transport paths, red crosses are targets; LAFM yields shorter, less entangled trajectories.

Architectural Insights and Practical Implications

LAFM’s architectural innovation lies in reparameterizing the source distribution via learned latent action priors, harmonizing initialization with coarse action intent, and thus structurally aligning denoising trajectories. Practically, this enables robots to learn smoother, more fluid behaviors from fragmented and noisy human demonstrations, with substantial improvements in real-world manipulation efficacy. The method achieves strong results with modest model size and no reliance on extensive pretraining, underscoring its scalability.

Theoretical Implications and Future Directions

By directly shaping the generative geometry via latent actions, LAFM bridges latent representation learning and structured generative modeling. The explicit prior adaptation provides a principled mechanism for tackling multi-modal action spaces and heteroscedastic distributions in imitation learning. The theoretical analysis establishes a formal link between source distribution initialization and policy transport efficiency, with potential extensions to probabilistic optimal transport and multi-agent settings.

Limitations include dependence on video-based LAMs, which may capture spurious environmental dynamics. Future refinements should focus on enhancing agent-specific disentanglement, potentially via stronger bottleneck constraints or hierarchical latent action modeling. Integrating adaptive flow matching priors into large-scale VLA pretraining pipelines represents a compelling avenue for general-purpose robotic foundation models.

Conclusion

LAFM advances generative denoising policy learning by structuring the source distribution through adaptive, latent action-guided priors. Empirical results demonstrate substantial gains in both simulated and real-world robotic manipulation, with performance exceeding that of larger, extensively pre-trained baselines. The approach not only improves control efficiency and generalization but also lays theoretical groundwork for future developments in structured generative modeling and scalable robotic learning.

Figure 7: Examples of real-world manipulation tasks used in evaluation—LAFM demonstrates robust performance across diverse settings.