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

Yanchen Huo; Ziying Song; Yadan Luo

cs.LGAug 29, 2026

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

Authors: Yanchen Huo, Ziying Song, Yadan Luo

Organizations: Nanyang Technological University · The University of Queensland

Abstract

Joint-Embedding Predictive Architectures (JEPAs) provide a powerful framework for latent world modeling and planning in a reconstruction-free manner. Although numerous JEPA-based approaches have been proposed to mitigate representation collapse, our experiments on localized, out-of-distribution visual noise reveal that performance degradation remains pronounced and unresolved. We propose Flow-JEPA (F-JEPA), a flow-based latent 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 switching from pointwise transition regression to stochastic trajectory-level prediction. F-JEPA raises mean success from $86\%$ to $92\%$ under clean observations and from $67\%$ to $86\%$ under noisy conditions. Further evaluations over varying perturbation severity and inference settings show that the performance advantage persists across a broad range of conditions. These results suggest that conditional flow matching provides a promising alternative to deterministic autoregressive prediction as a dynamics formulation in JEPA world models.

Figures & tables

Figure 1: Latent-dynamics formulation affects planning robustness. F-JEPA, which instantiates trajectory-level vector-field dynamics, achieves higher average success under clean observations and substantially smaller degradation under noisy observations than the autoregressive LeWM baseline across four environments. The

y

-axis starts at 0.5.

Figure 3: Overview of F-JEPA. (a) training pipeline. The conditional flow matching objective trains the predictor to transport Gaussian source tokens toward future latent representations conditioned on the current latent state and action sequence, while SIGReg regularizes the visual encoder to prevent representation collapse. (b) predictor architecture. The predictor jointly updates future latent tokens using Transformer blocks with flow-time conditioning and action-dependent context.

Method	Two-Room	Reacher	Push-T	OGBench-Cube	Mean
PLDM	97	78	78	65	80
DINO-WM	100	79	74	86	85
LeWM	87	86	96	74	86
F-JEPA	100	90	96	82	92

Table 1: Planning success rates (%) under clean visual conditions. F-JEPA achieves the highest mean success rate (92%) and the best or tied-best performance on three of the four environments, showing that the alternative dynamics formulation preserves strong clean-condition planning performance. Results for PLDM ( Sobal et al., 2026 ) , DINO-WM ( Zhou et al., 2024 ) , and LeWM are taken from ( Maes et al., 2026 ) . The best result in each column is shown in bold.

Method	Two-Room	Reacher	Push-T	OGBench-Cube	Mean
LeWM	76	74	58	58	67
F-JEPA	96	90	90	66	86
$\Delta$	+20	+16	+32	+8	+19

Table 2: Planning success rates (%) under noisy visual conditions. F-JEPA improves performance across all four environments and increases the mean success rate from 67% to 86%. The final row reports the absolute percentage-point difference from LeWM.

Figure 4: Qualitative rollouts in Two-Room. Both LeWM (left) and F-JEPA (right) reach the goal under clean observations, whereas LeWM fails after task-irrelevant background noise is introduced and F-JEPA remains successful, illustrating improved robustness to visual perturbations.

Figure 5: Qualitative rollouts in OGBench-Cube. Both LeWM (left) and F-JEPA (right) complete the task under clean observations, whereas LeWM fails after task-irrelevant background noise is introduced and F-JEPA remains successful, illustrating improved robustness to visual perturbations.

	Gaussian-patch radius $r$ (pixels)
Method	10	15	20	25	30	35	40
LeWM	80	84	72	74	68	74	72
F-JEPA	88	86	82	88	82	90	84
$\Delta$	+8	+2	+10	+14	+14	+16	+12

Table 3: Reacher success rates (%) across Gaussian-patch radii. F-JEPA consistently outperforms LeWM across all tested perturbation radii, indicating that its robustness gain is not specific to a particular spatial noise scale. The noise standard deviation is fixed at

s=100

, and

\Delta

denotes the absolute percentage-point improvement over LeWM.

	Noise standard deviation $s$
Method	50	100	150	200
LeWM	82	74	66	56
F-JEPA	88	90	78	76
$\Delta$	+6	+16	+12	+20

Table 4: Reacher success rates (%) across Gaussian-noise standard deviations. F-JEPA consistently outperforms LeWM across all tested noise levels and maintains an advantage even at the largest perturbation magnitude. The patch radius is fixed at

r=35

, and

\Delta

denotes the absolute percentage-point improvement over LeWM.

	Noise standard deviation $s$
Method	50	100	150	200
LeWM	82	74	66	56
F-JEPA	88	90	78	76
$\Delta$	+6	+16	+12	+20

r=35

, and

\Delta

denotes the absolute percentage-point improvement over LeWM.

	Clean		Noisy
Flow steps	Success	Time	Success	Time
4	96	31.7	82	35.0
8	96	42.7	90	47.7
16	96	66.9	82	78.0
32	92	111.8	84	130.2

Table 5: Push-T success rates (%) across flow integration steps and corresponding inference cost (s). 8 Euler steps achieve the highest success rates. Times report wall-clock seconds over 50 rollouts; bold denotes the highest success rate within each condition.

	Clean			Noisy
Method	10	20	30	10	20	30
LeWM	94	92	96	60	56	58
F-JEPA	96	96	96	90	88	84

Table 6: Push-T success rates (%) across CEM iterations. F-JEPA maintains strong clean performance across all tested planning budgets and achieves its best noisy success rate with only 10 CEM iterations. Under the same 10-iteration budget, F-JEPA substantially outperforms LeWM under noise. Bold denotes the highest success rate for each method and visual condition.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Test environments. Representative observations from the four evaluation environments: Two-Room, Reacher, Push-T, and OGBench-Cube, covering navigation, continuous control, and object manipulation tasks.

Figure 7: Qualitative rollouts in Reacher. LeWM (left) and F-JEPA (right) are shown under clean and noisy observations. Under clean observations, both methods complete the task, while F-JEPA reaches the target with higher precision. After visual perturbations are introduced, LeWM fails to reach the target, whereas F-JEPA remains successful.

Figure 8: Qualitative rollouts in Push-T. LeWM (left) and F-JEPA (right) are shown under clean and noisy observations. Both methods complete the task under clean observations, whereas F-JEPA remains successful under visual perturbations while LeWM fails.

		Gaussian-patch radius $r$ (pixels)
Environment	Method	10	15	20	25	30	35	40
Two-Room	LeWM	78	76	80	82	78	76	72
	F-JEPA	100	100	98	98	92	96	80
	$\Delta$	+22	+24	+18	+16	+14	+20	+8
Push-T	LeWM	58	42	32	22	6	8	4
	F-JEPA	90	80	40	26	22	8	14

Appendix

Table 7: Additional analysis across Gaussian-patch radii. Complementing the Reacher results in the main text, F-JEPA matches or outperforms LeWM across all tested radii in Two-Room, Push-T, and OGBench-Cube, indicating that its robustness advantage persists across different spatial perturbation scales. The noise standard deviation is fixed at

s=100

, and

\Delta

denotes the absolute percentage-point difference between F-JEPA and LeWM.

Figure 9: Success-rate trends across Gaussian-patch radii. F-JEPA generally maintains higher success rates than LeWM across different spatial perturbation scales in all four environments, complementing the quantitative results in Tables 3 and 7 .

		Noise standard deviation $s$
Environment	Method	50	100	150	200
Two-Room	LeWM	76	76	70	68
	F-JEPA	100	96	84	80
	$\Delta$	+24	+20	+14	+12
Push-T	LeWM	86	58	56	54
	F-JEPA	94	90	72	72

Appendix

Table 8: Additional analysis across Gaussian-noise magnitudes. Across Two-Room, Push-T, and OGBench-Cube, F-JEPA generally maintains higher success than LeWM as the perturbation magnitude increases, showing that the observed robustness difference is not confined to the noise level used in the main evaluation. The patch radius is fixed at

r=35

for Two-Room and OGBench-Cube and at

r=10

for Push-T, and

\Delta

denotes the absolute percentage-point difference between F-JEPA and LeWM.

Figure 10: Success-rate trends across Gaussian-noise magnitudes. F-JEPA generally retains a performance advantage over LeWM as the noise standard deviation increases, indicating improved robustness across a broad range of perturbation strengths. Corresponding numerical results are reported in Tables 5 and 8 .

Figure 11: Planning performance and inference time across flow integration steps. Success rates under clean and noisy observations and corresponding planning times are shown across different Euler integration budgets for all four environments. Increasing the number of flow steps does not consistently improve planning performance, while inference time grows substantially with additional integration steps. Overall, 8 steps provide a favorable performance-efficiency trade-off across the evaluated tasks. Corresponding numerical results are reported in Tables 5 and 9 .

Environment	Flow steps	Success (%)	Time (s)	Success (%)	Time (s)
		Clean		Noisy
Two-Room	4	100	22.5	96	32.2
	8	100	34.4	96	44.5
	16	100	52.3	94	66.5
	32	98	98.2	96	109.5
Reacher	4	84	47.1	86	47.0

Appendix

Table 9: Additional analysis across flow integration steps. Complementing the Push-T results in the main text, 8 Euler steps achieve the best or tied-best planning performance across the remaining environments, while additional integration steps substantially increase runtime without yielding consistent performance gains. Times report wall-clock seconds over 50 rollouts.

Method	Inference latency (ms)	Planning time (ms)	50 rollouts (s)
LeWM	2.1 / prediction	105.0	21.7
F-JEPA	5.0 / flow step	400.0	34.4

Appendix

Table 10: Inference-time comparison between LeWM and F-JEPA. F-JEPA incurs higher dynamics and per-planning latency due to iterative flow integration, while the end-to-end runtime gap over 50 rollouts is smaller; over these rollouts, LeWM requires substantially more replanning during task execution. Both methods use 10 CEM refinement iterations.

Explore similar work

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling
Aug 7, 2026An Lanji, Dawei Liu, Jin Li +3Video Joint Embedding Predictive Architecture Joint-Embedding Predictive Architectures
University of Electronic Science and Technology of China, Chengdu, China.
PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Latent World Models Joint-Embedding Predictive Architectures
School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China
A Generalization Theory for JEPA-Based World Models
Jun 25, 2026Jingyi Cui, Qi Zhang, Hongwei Wen +1Joint-Embedding Predictive Architectures World Models
State Key Lab of General AI, School of Intelligence Science and Technology, Peking University · University of Sydney

Aug 7, 2026, cs.CVJ/K move · Enter open · S save

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

An Lanji, Dawei Liu, Jin Li, Haoran Xu+2

University of Electronic Science and Technology of China, Chengdu, China.

Video Joint Embedding Predictive Architecture Joint-Embedding Predictive Architectures Video World Models World Latent Space Action-Conditioned Planning

Open paper AskPDF arXiv

Joint-Embedding Predictive Architectures (JEPAs) have emerged as a principled framework for self-supervised learning of world models in compact latent spaces, yet existing methods are fragmented: some predict masked parts of a single image in latent space (I-JEPA), others learn to predict global photometric transformations (Image World Models), while video-scale JEPAs predict future temporal states and are post-trained for action-conditioned planning (V-JEPA~2, DINO-World, DINO-WM). These objectives are treated as distinct recipes with separate encoders, predictors, and anti-collapse regularizers, hindering a single model from unifying image-level and video-level world modeling. We present UniJEPA, a unified JEPA that jointly learns photometric prediction (image-level transformations) and temporal prediction (video-level next-state dynamics) in one shared latent space. A single end-to-end objective, composed of a next-embedding prediction loss and a Gaussian regularizer, yields a provably anti-collapse encoder-predictor pair trainable from raw pixels without EMA, stop-gradient, or pre-trained encoders. We show that the same latent space supports controllable abstraction: photometric prediction learns invariant structure while temporal prediction learns equivariant dynamics. After action-conditioned post-training on offline trajectories, UniJEPA enables zero-shot planning by treating goal features as prediction targets. On image, video, and control benchmarks, UniJEPA matches or surpasses task-specific JEPAs while requiring a single loss hyperparameter, and plans up to tens of times faster than generative world models at comparable accuracy.

Abstract

Explore similar work

UniJEPA: A Unified Joint-Embedding Predictive Architecture for Task-Agnostic Visual World Modeling

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

A Generalization Theory for JEPA-Based World Models