PaperlayerView exact revision on arXiv

Source-linked AI summary

Flow-JEPA: Flow Matching for Robust Latent Dynamics in JEPA World Models

Yanchen Huo, Ziying Song, Yadan Luo

arXiv:2608.29029v1cs.LGcs.AI

TL;DR

LeWM’s deterministic autoregressive latent predictor can accumulate errors and remain sensitive to task-irrelevant visual perturbations. Flow-JEPA replaces it with conditional flow matching over jointly generated future latent trajectories, improving mean success from 86% to 92% under clean observations and from 67% to 86% under noisy conditions.

  • Problem

    LeWM’s repeated one-step transitions can accumulate prediction errors, while deterministic point-wise regression does not explicitly handle latent perturbations from task-irrelevant visual changes.

  • Method

    Flow-JEPA uses conditional flow matching to transport a Gaussian source trajectory toward complete future latent trajectories conditioned on current observations and actions.

  • Results

    Mean success rises from 86% to 92% under clean observations and from 67% to 86% under noisy conditions.

  • Takeaways & Limitations

    Conditional flow matching is a promising alternative to deterministic autoregressive latent dynamics for robust JEPA world models.

  • Takeaways & Limitations

    Experiments focus on relatively short-horizon control tasks and primarily localized Gaussian visual perturbations, while inference requires multiple vector-field evolutions.

Abstract

from arXiv · show

Joint-Embedding Predictive Architectures (JEPAs) have shown strong potential for learning compact predictive representations, and LeWorldModel (LeWM) extends this paradigm to reconstruction-free latent world modeling from pixels. However, its deterministic autoregressive predictor generates future states through repeated one-step transitions, which can accumulate errors and remain sensitive to task-irrelevant visual perturbations. In this work, we propose Flow-JEPA (F-JEPA), a conditional flow matching dynamics model that jointly generates a sequence of future latent states conditioned on the current observation and actions. A Gaussian distribution serves as the flow source, exposing the vector field to perturbed latent trajectories as it learns to transport them toward clean future representations. This formulation retains the reconstruction-free JEPA framework while replacing point-wise transition regression with stochastic trajectory-level prediction. F-JEPA raises mean success from $86\%$ to $92\%$ under clean observations and from $67\%$ to $86\%$ under noisy conditions, suggesting that conditional flow matching provides a promising alternative to deterministic autoregressive dynamics in JEPA world models.

1 INTRODUCTION

Flow-JEPA replaces LeWM’s deterministic autoregressive latent dynamics with conditional flow matching over complete future trajectories. Across four environments, it improves clean planning success and reduces degradation under noisy visual observations.

  • Background: LeWorldModel extends reconstruction-free JEPA modeling to predict future latent representations from pixel observations and actions.JEPAs forecast embeddings rather than reconstructing pixels, while LeWM applies this framework to reward-free world modeling.
  • Problem: LeWM’s dynamics predictor has two limitations: autoregressive error accumulation and deterministic point-wise alignment.Repeated one-step transitions propagate errors, while point-wise regression does not explicitly model recovery from local latent perturbations.
  • Problem: Task-irrelevant visual perturbations can make JEPA-style learned dynamics sensitive despite leaving the underlying environment dynamics unchanged.
  • Approach: F-JEPA transports a Gaussian source trajectory toward conditional future latent trajectories, jointly predicting horizons instead of recursively applying one-step transitions.This replaces point-wise latent regression with stochastic conditional transport over complete action-conditioned trajectories.
  • Results: F-JEPA improves average planning performance under clean observations and substantially reduces performance degradation under visual perturbations.The introduction reports gains across four environments under both clean and noisy conditions.

2 RELATED WORK

Related work develops latent world models, reconstruction-free JEPA predictors, and flow matching for conditional generation. Flow-JEPA combines these lines by applying conditional flow matching to future JEPA latent states.

  • Latent World Models: Latent world models compress high-dimensional observations into compact states and learn predictive dynamics for imagined control and policy learning.
  • Joint-Embedding Predictive Architectures: JEPAs forecast target embeddings without reconstructing raw inputs, and recent systems apply this principle to end-to-end world models for control.Examples include I-JEPA, V-JEPA, DINO-WM, PLDM, and LeWorldModel.
  • Flow Matching for Conditional Generation: Flow matching learns a time-dependent vector field that transports samples from a source noise distribution to a target data distribution.The framework supports flexible conditional distributions and has been used for image, video, and robot action generation.
  • Flow Matching for Conditional Generation: Flow-JEPA applies conditional flow matching to future JEPA latent states as a generative alternative to deterministic autoregressive latent dynamics.

3 RETHINKING LATENT DYNAMICS IN LEWORLDMODEL

LeWorldModel learns action-conditioned latent dynamics with a point-wise next-embedding objective and SIGReg regularization. Its recursive transitions can accumulate errors, while deterministic alignment does not constrain behavior around perturbed latent states.

  • JEPA World Modeling: An offline observation-action dataset is encoded into compact latent states, and an action-conditioned predictor forecasts the next latent state.The visual encoder maps pixel observations to latent representations, while the dynamics predictor uses actions to predict transitions.
  • Representation Collapse: The original LeWM training objective uses point-wise next-embedding prediction, which can admit a constant collapsed representation with zero prediction error.Such collapse fails to preserve information about environment dynamics.
  • Representation Collapse: SIGReg regularizes encoder embeddings by projecting them onto sampled directions and applying univariate Epps-Pulley tests.Its objective encourages approximately isotropic Gaussian embeddings and prevents mapping all inputs to one constant.
  • Robust Predictive Representations: JEPA reconstruction-free prediction aims to abstract away task-irrelevant appearance, but visual perturbations can still shift latent representations when physical state is unchanged.
  • Recursive Dynamics: LeWM models multi-step dynamics through recursive one-step predictions, so errors at one step alter later inputs and propagate through the rollout.This is especially problematic when planning depends on several imagined transitions.
  • Perturbed Latent States: Deterministic point-wise alignment supervises clean target embeddings but does not explicitly constrain dynamics under latent uncertainty caused by visual perturbations.

4 METHOD

F-JEPA formulates multi-step latent dynamics as conditional flow matching, jointly evolving future latent trajectories from Gaussian source samples toward action-conditioned future representations. Its Transformer predictor uses flow-time conditioning, while SIGReg regularizes the visual representation space.

  • Conditional flow matching: F-JEPA learns a continuous-time vector field that transports Gaussian source trajectories toward conditional future latent representations.The predictor estimates trajectory velocity at randomly sampled intermediate flow states.
  • Data and conditioning: Future latent tokens are processed jointly by Transformer blocks conditioned on the current latent state, encoded actions, temporal positions, and flow time.Flow time enters through sinusoidal embeddings and adaptive layer normalization, while cross-attention uses the current latent and action sequence as context.
  • Conditional flow matching: The entire future sequence is modeled as one joint variable, allowing prediction horizons to evolve jointly instead of recursively feeding intermediate predictions forward.Training avoids numerical ODE integration and does not require autoregressive intermediate transitions.
  • Data and conditioning: Training samples contain a current observation, h future observations, and h temporally aligned action blocks for a prediction horizon of h.Each action block contains five consecutive low-level environment actions.
  • Optimization: SIGReg regularizes embeddings of current and future observations to prevent representation collapse while the flow objective learns action-conditioned latent dynamics.The regularizer constrains representation geometry, whereas flow matching learns dynamics along stochastic source-target paths.
  • Inference: At inference, Euler integration transforms a Gaussian source trajectory into the predicted future latent trajectory using the learned conditional vector field.With N flow steps, the integration interval is Δτ = 1/N and all future latent states are updated jointly at each integration step.

5 EXPERIMENTS

Experiments evaluate F-JEPA across four environments under clean and visually perturbed observations, including robustness, integration-step, and planning-budget studies. F-JEPA achieves higher mean success than LeWM and maintains its advantage across perturbation severities and matched planning settings.

  • Planning performance and robustness: F-JEPA matches or exceeds LeWM across all four environments under clean observations, raising mean success from 86% to 92%.The clean-condition comparison reports the highest mean success rate for F-JEPA.
  • Planning performance and robustness: Under visual perturbations, F-JEPA raises mean success from 67% to 86% and reduces mean degradation from 19 points for LeWM to 6 points.The noisy-condition improvement is reported across the evaluated environments.
  • Planning performance and robustness: Qualitative rollouts show both methods succeeding under clean observations, while F-JEPA remains successful under background noise in Two-Room and OGBench-Cube where LeWM fails.These examples visualize robustness to task-irrelevant visual perturbations.
  • Ablation studies: F-JEPA consistently outperforms LeWM across tested Reacher perturbation radii and noise levels, indicating robustness beyond a single spatial-noise configuration.The radius study fixes s = 100, while the noise-level study fixes r = 35.
  • Ablation studies: Eight Euler steps preserve the best clean performance and achieve the highest noisy success, whereas additional flow steps increase runtime without improving planning performance.The study uses 8 flow steps in the main experiments.
  • Ablation studies: F-JEPA maintains 96% clean success across the tested CEM budgets and achieves its highest noisy success with 10 iterations, including under matched budgets.The matched-budget study compares F-JEPA with the reported LeWM baseline using the same iteration count.

6 CONCLUSION

F-JEPA replaces deterministic one-step latent prediction with stochastic transport over complete future trajectories in reconstruction-free JEPA world models. Across four control environments, it improves clean-observation planning and reduces degradation from task-irrelevant visual perturbations, while small-step integration remains sufficient in the studied tasks.

  • Conclusion: F-JEPA jointly transports complete future latent trajectories instead of autoregressively predicting one step at a time.The model uses conditional flow matching for stochastic trajectory prediction, allowing all prediction horizons to evolve jointly.
  • Conclusion: F-JEPA improves average planning performance under clean observations and substantially reduces degradation under visual perturbations across four control environments.These results support conditional flow matching as a promising alternative to deterministic autoregressive latent dynamics.
  • Conclusion: Robustness persists across perturbation strengths without requiring increasingly fine ODE integration or a larger CEM planning budget.The conclusion reports that these ablations preserve the robustness improvement while avoiding larger inference-time budgets.
  • Limitations and future work: The experiments focus on relatively short-horizon control tasks and primarily localized Gaussian visual perturbations.The authors identify longer-horizon tasks and broader distribution shifts as future evaluation directions.
  • Limitations and future work: Flow-based prediction requires multiple vector-field evolutions at inference time, making prediction efficiency an important remaining direction.The authors report that a small number of integration steps is sufficient in the studied tasks but still identify efficiency as future work.

A IMPLEMENTATION DETAILS

The implementation uses Gaussian flow sources and configurable attention over future trajectory tokens, with a Transformer-based predictor and environment-specific training choices. Models share a compact visual-action representation and are trained with a five-step prediction horizon.

  • Flow source and attention configuration: The implementation considers standard-noise Gaussian sources and either bidirectional or causal self-attention among future trajectory tokens.Action conditioning remains temporally causal in both attention variants.
  • Architecture: All observations are resized to 224 × 224, encoded with a ViT-Tiny, and paired with an MLP action encoder using 192-dimensional projections.The flow predictor uses six Transformer blocks with 16 attention heads and an MLP hidden dimension of 2048.
  • Training details: F-JEPA models are trained for 20 epochs with batch size 128, prediction horizon 5, and history size 1.Push-T freezes LeWM-initialized visual components, while Reacher keeps them trainable as a warm start.
  • Hardware: All training and evaluation experiments use a single NVIDIA RTX 6000 Ada GPU.

B.1 TEST ENVIRONMENTS AND DATASETS

The evaluation covers four visually distinct control environments spanning navigation, continuous control, and manipulation. Following the LeWM dataset construction, each latent transition is conditioned on an action block of five low-level actions rather than one action.

  • Test environments: The four environments are Two-Room, Reacher, Push-T, and OGBench-Cube.They cover goal-directed navigation, planar reaching, planar object pushing, and goal-conditioned cube manipulation.
  • Datasets: Observations are separated by a frame skip of 5, with the intervening low-level actions concatenated into an action block.This follows the LeWM dataset construction.
  • Datasets: Each latent transition is conditioned on one five-action block rather than a single low-level action.

B.2 ADDITIONAL RESULTS

Additional results show that F-JEPA’s robustness advantage persists across perturbation scales and environments, while eight Euler steps offer a favorable accuracy-efficiency trade-off. Flow integration increases per-cycle cost, but the end-to-end runtime gap is smaller because F-JEPA often requires fewer replanning cycles.

  • Qualitative rollouts: F-JEPA remains more reliable than LeWM after visual perturbations in qualitative Reacher and Push-T rollouts.Both methods succeed under clean observations, while F-JEPA remains successful under perturbations and reaches the Reacher target with higher precision.
  • Robustness ablations: F-JEPA generally maintains higher success rates than LeWM across Gaussian-patch radii and noise magnitudes.The trends span all four environments and complement the numerical ablations.
  • Robustness ablations: F-JEPA matches or exceeds LeWM across every tested patch radius in Two-Room, Push-T, and OGBench-Cube.The patch-radius sweep fixes the noise standard deviation at s = 100.
  • Robustness ablations: F-JEPA remains stronger across nearly all tested noise magnitudes, including the largest perturbations, with a small OGBench-Cube reversal at s = 150.Push-T remains advantaged or tied throughout the spatial-radius sweep despite increasing difficulty.
  • Flow-step ablations: Eight Euler steps provide a favorable accuracy-efficiency trade-off because additional steps increase runtime without consistently improving planning performance.Two-Room is near saturation with 4–8 steps, while Reacher and OGBench-Cube achieve their strongest clean and noisy results with 8 steps.
  • Inference-time analysis: F-JEPA is approximately 3.8× slower per planning cycle but only 1.6× slower over 50 rollouts than LeWM.The smaller end-to-end gap arises because LeWM requires substantially more replanning cycles during complete rollouts.
  • Inference-time analysis: F-JEPA’s joint trajectory updates may become relatively more favorable at longer horizons than autoregressive rollout.Autoregressive rollout requires an increasing number of sequential model evaluations as the prediction horizon grows.
Loading 2608.29029v1…