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

# Flow-JEPA: Flow Matching for Robust Latent Dynamics in JEPA World Models
Source: [https://arxiv.org/html/2608.29029](https://arxiv.org/html/2608.29029)
Ziying Song††thanks:Corresponding authorAffiliation:Nanyang Technological UniversityYadan LuoAffiliation:The University of Queensland

###### Abstract

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 from86%86\\%to92%92\\%under clean observations and from67%67\\%to86%86\\%under noisy conditions, suggesting that conditional flow matching provides a promising alternative to deterministic autoregressive dynamics in JEPA world models\. Code is available at[https://github\.com/HuoYanchen/Flow\-JEPA](https://github.com/HuoYanchen/Flow-JEPA)\.

## 1Introduction

World models have emerged as a promising type of deep neural networks for learning predictive representations of environment dynamics\. By modeling the future consequences of actions, world models provide a foundation for planning, policy learning, and synthetic data generation\. Early and influential approaches learn compact latent dynamics from high\-dimensional observations and use the learned model for imagination\-based control or policy optimization\([Ha and Schmidhuber, 2018](https://arxiv.org/html/2608.29029#bib.bib1);[Hafner et al\., 2019a](https://arxiv.org/html/2608.29029#bib.bib2)\)\. More recent systems, such as IRIS\([Micheli et al\., 2022](https://arxiv.org/html/2608.29029#bib.bib3)\), further demonstrate that sequence models can serve as powerful world models by combining learned visual tokenization with autoregressive dynamics prediction\. Despite their success, many world models rely on pixel\-level reconstruction or generative prediction objectives\. While such objectives preserve rich visual details, they can also force the model to spend capacity on task\-irrelevant factors such as background or lighting\. This motivates reconstruction\-free latent world models that predict future states in an abstract representation space rather than reconstructing raw pixels\.

Joint\-Embedding Predictive Architectures \(JEPAs\)\([LeCun, 2022](https://arxiv.org/html/2608.29029#bib.bib24)\)provide a fundamental framework for such reconstruction\-free modeling\. Instead of predicting pixels, JEPA\-style methods encode observations into latent embeddings and train a predictor to forecast missing or future representations\. This idea has recently been extended to world modeling\. LeWorldModel \(LeWM\)\([Maes et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib7)\)trains a compact end\-to\-end JEPA world model directly from pixels in reward\-free environments\. It learns an encoder that maps observations to latent states and an action\-conditioned predictor that forecasts future latent embeddings, while using Sketched\-Isotropic\-Gaussian Regularizer \(SIGReg\)\([Balestriero and LeCun, 2025](https://arxiv.org/html/2608.29029#bib.bib8)\)to prevent representation collapse without reconstruction losses, frozen encoders, exponential moving averages, or auxiliary supervision\.

Figure 1:Flow\-JEPA \(F\-JEPA\) improves both planning performance and robustness to visual perturbations\.F\-JEPA replaces the deterministic autoregressive dynamics of LeWM with conditional flow matching over future latent trajectories\. Across four environments, it achieves higher average success under clean conditions and substantially reduces performance degradation under noisy observations\. Theyy\-axis starts at 0\.5\.However, the dynamics predictor in LeWM retains two major limitations:

- •Error accumulation\.The future latent states are generated autoregressively through repeated one\-step transitions, causing prediction errors to propagate and accumulate over the rollout horizon\.
- •Deterministic point\-wise alignment\.The predictor is trained through deterministic point\-wise regression between predicted and target embeddings\. While this objective encourages accurate prediction on clean training trajectories, it provides no explicit mechanism for modeling or recovering from local perturbations of the latent state\.

Consequently, although JEPA\-style world models avoid pixel reconstruction and are intended to capture abstract predictive representations, their learned dynamics can remain sensitive to task\-irrelevant visual perturbations\.

To address these limitations, we propose Flow\-JEPA \(F\-JEPA\), a flow matching latent dynamics model for JEPA\-based world modeling\. Flow matching learns a time\-dependent vector field that transports samples from a source distribution toward a conditional target distribution and has shown strong performance in continuous generation problems such as image, video, and robot action generation\([Lipman et al\., 2022](https://arxiv.org/html/2608.29029#bib.bib9);[Wan et al\., 2025](https://arxiv.org/html/2608.29029#bib.bib11);[Black et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib10)\)\. In this work, we use it to model future latent trajectories\. Given the current latent state and an action sequence, F\-JEPA transports a Gaussian stochastic source trajectory toward the corresponding future latent trajectory\. This removes the recursive dependence between predicted future states while replacing point\-wise latent regression with a flexible flow matching objective\. The Gaussian source used in flow matching is naturally compatible with the latent geometry promoted by SIGReg, while the two mechanisms also operate at complementary stages of the model\. SIGReg prevents representation collapse and regularizes the encoder outputs, whereas flow matching trains the dynamics predictor over Gaussian stochastic neighborhoods of the learned latent states\. We hypothesize that learning transport from these perturbed trajectories promotes local stability of the latent dynamics and improves robustness to perturbation\-induced shifts in visual representations\.

Our experiments show that this reformulation improves average planning performance and robustness\. We demonstrate that LeWM can be substantially affected by task\-irrelevant visual perturbations, although they leave the underlying environment dynamics unchanged\. F\-JEPA improves average performance under clean conditions and exhibits substantially smaller degradation under perturbed conditions, suggesting that flow\-based latent trajectory prediction provides a more robust alternative to deterministic autoregressive dynamics in JEPA\-style world models\.

Our main contributions are summarized as follows:

- •We empirically demonstrate that LeWorldModel is sensitive to task\-irrelevant visual perturbations during planning, despite these perturbations leaving the underlying environment dynamics unchanged\.
- •We propose F\-JEPA, a flow matching JEPA world model that replaces deterministic autoregressive transitions with stochastic conditional transport of complete action\-conditioned future latent trajectories\.
- •We show that F\-JEPA improves average planning performance under clean conditions while substantially increasing robustness to visual perturbations\.

## 2Related work

#### Latent World Models\.

Latent world models compress high\-dimensional observations into compact state representations and learn predictive dynamics in the resulting latent space\. Early approaches predominantly adopt recurrent architectures and task\-specific reward supervision\([Ha and Schmidhuber, 2018](https://arxiv.org/html/2608.29029#bib.bib1);[Hafner et al\., 2019b](https://arxiv.org/html/2608.29029#bib.bib12)\)\. The Dreamer family replaces online action search with actor\-critic learning over trajectories imagined by the world model\([Hafner et al\., 2019a](https://arxiv.org/html/2608.29029#bib.bib2);[Hafner et al\., 2020](https://arxiv.org/html/2608.29029#bib.bib13);[Hafner et al\., 2023](https://arxiv.org/html/2608.29029#bib.bib14);[Hafner et al\., 2025](https://arxiv.org/html/2608.29029#bib.bib15)\)\. Other generative approaches include IRIS\([Micheli et al\., 2022](https://arxiv.org/html/2608.29029#bib.bib3)\), which models discrete visual tokens with an autoregressive Transformer, and DIAMOND\([Alonso et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib16)\), which uses diffusion\-based observation prediction for Atari control\. Recent work has further explored stochastic generative forecasting directly in pretrained feature spaces\([Walker et al\., 2025](https://arxiv.org/html/2608.29029#bib.bib31);[Porcher et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib30)\)\. Our work adopts a reward\-free and reconstruction\-free JEPA setting, where action\-conditioned dynamics are learned directly in representation space\.

#### Joint\-Embedding Predictive Architectures\.

Joint\-Embedding Predictive Architectures \(JEPAs\) learn to forecast target embeddings without reconstructing raw inputs\. I\-JEPA\([Assran et al\., 2023](https://arxiv.org/html/2608.29029#bib.bib4)\)and V\-JEPA\([Bardes et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib5);[Assran et al\., 2025](https://arxiv.org/html/2608.29029#bib.bib27)\)show that representation prediction can learn strong image and video features without pixel reconstruction\. This idea has recently been adopted for world modeling\. DINO\-WM\([Zhou et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib6)\)and PLDM\([Sobal et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib18)\)are important early examples of end\-to\-end JEPA\-style world models for control\. Our work is based on LeWorldModel\([Maes et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib7)\), which trains a compact end\-to\-end JEPA world model using a next\-embedding prediction objective and SIGReg\([Balestriero and LeCun, 2025](https://arxiv.org/html/2608.29029#bib.bib8)\)\. Our method retains the Gaussian\-regularized encoder while replacing the autoregressive one\-step predictor with conditional flow matching over complete future latent trajectories\.

#### Flow Matching for Conditional Generation\.

Flow matching is a generative modeling framework that learns a time\-dependent vector field transporting samples from a source noise distribution to a target data distribution\.\([Lipman et al\., 2022](https://arxiv.org/html/2608.29029#bib.bib9);[Liu, 2022](https://arxiv.org/html/2608.29029#bib.bib17);[Liu et al\., 2022](https://arxiv.org/html/2608.29029#bib.bib25)\)Compared with direct regression, flow\-based objectives provide a flexible way to model continuous conditional distributions and have been applied to domains such as image generation, video generation, and robot action generation\([Wan et al\., 2025](https://arxiv.org/html/2608.29029#bib.bib11);[Black et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib10)\)\. Recent work has explored replacing uninformed Gaussian initialization with structured source distributions\([Jia et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib19);[Kim et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib26)\), exploiting temporal continuity and strengthening conditional generation\. Our work applies conditional flow matching to future JEPA latent states, providing a generative alternative to deterministic autoregressive latent dynamics\.

## 3Rethinking Latent Dynamics in LeWorldModel

### 3\.1JEPA World Modeling and Representation Collapse

We consider an offline dataset of observation–action trajectories𝒟=\{\(ot,at\)\}t=1T\\mathcal\{D\}=\\\{\(o\_\{t\},a\_\{t\}\)\\\}\_\{t=1\}^\{T\}, whereoto\_\{t\}denotes a pixel observation andata\_\{t\}denotes the corresponding action\. A visual encoderfθf\_\{\\theta\}maps each observation into a compact latent representationztz\_\{t\}, and an action\-conditioned dynamics predictorgϕg\_\{\\phi\}predicts the next latent state:

z^t\+1=gϕ​\(zt,at\)\.\\hat\{z\}\_\{t\+1\}=g\_\{\\phi\}\(z\_\{t\},a\_\{t\}\)\.\(1\)The encoder and predictor are trained jointly using a point\-wise next\-embedding objective,

ℒpred=‖z^t\+1−zt\+1‖22\.\\mathcal\{L\}\_\{\\mathrm\{pred\}\}=\\left\\\|\\hat\{z\}\_\{t\+1\}\-z\_\{t\+1\}\\right\\\|\_\{2\}^\{2\}\.\(2\)
When optimized alone, this objective admits a trivial collapsed solution\. In particular, the encoder may map every observation to the same constant representationzt=cz\_\{t\}=c, while the predictor outputs the same constant independently of its inputs\. This yields zero prediction error without preserving any information about the environment dynamics\. Representation collapse is therefore a central challenge in end\-to\-end JEPA training\.

To prevent this trivial solution, LeWM\([Maes et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib7)\)applies the SIGReg\([Balestriero and LeCun, 2025](https://arxiv.org/html/2608.29029#bib.bib8)\)to the encoder embeddings\. LetZ∈ℝN×dZ\\in\\mathbb\{R\}^\{N\\times d\}denote the latent embeddings collected across observations and trajectories\. SIGReg samplesMMunit\-norm directionsu\(m\)∈𝕊d−1u^\{\(m\)\}\\in\\mathbb\{S\}^\{d\-1\}and projects the embeddings onto each direction:

h\(m\)=Z​u\(m\)\.h^\{\(m\)\}=Zu^\{\(m\)\}\.\(3\)A univariate Epps\-Pulley\([Epps and Pulley, 1983](https://arxiv.org/html/2608.29029#bib.bib23)\)test statistic𝒯\\mathcal\{T\}is then applied to each projection,

ℒSIG​\(Z\)=1M​∑m=1M𝒯⁡\(h\(m\)\)\.\\mathcal\{L\}\_\{\\mathrm\{SIG\}\}\(Z\)=\\frac\{1\}\{M\}\\sum\_\{m=1\}^\{M\}\\mathcal\{T\}\\\!\\left\(h^\{\(m\)\}\\right\)\.\(4\)By encouraging all one\-dimensional projections to follow a standard Gaussian distribution, SIGReg promotes an approximately isotropic Gaussian embedding distribution and prevents the encoder from mapping all inputs to a constant\. The original LeWM objective is therefore

ℒLeWM=ℒpred\+λSIG​ℒSIG\.\\mathcal\{L\}\_\{\\mathrm\{LeWM\}\}=\\mathcal\{L\}\_\{\\mathrm\{pred\}\}\+\\lambda\_\{\\mathrm\{SIG\}\}\\mathcal\{L\}\_\{\\mathrm\{SIG\}\}\.\(5\)

### 3\.2From Non\-Collapse to Robust Predictive Representations

A central motivation of JEPA\-style modeling is to predict in representation space rather than reconstructing every pixel\-level detail\. Ideally, observations that differ only in task\-irrelevant appearance should induce similar predictive dynamics\. In this sense, reconstruction\-free prediction provides an opportunity to abstract away nuisance visual information and retain features that are most relevant to the temporal evolution of the environment\([Sun et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib29)\)\. However, a perturbed observation can induce shifts in the latent representation even when the corresponding physical state is unchanged\. Robust world modeling consequently requires not only a well\-structured encoder space, but also a dynamics model that remains reliable under moderate deviations within that space\.

From this perspective, we identify two limitations in the original LeWM dynamics formulation\. First, LeWM models multi\-step dynamics through recursive one\-step predictions:

z^t\+k=gϕ\(z^t\+k−1,at\+k−1\),k=1,…,h,\\hat\{z\}\_\{t\+k\}=g\_\{\\phi\}\\left\(\\hat\{z\}\_\{t\+k\-1\},a\_\{t\+k\-1\}\\right\),\\qquad k=1,\\ldots,h,\(6\)withz^t=zt\\hat\{z\}\_\{t\}=z\_\{t\}\. Each predicted latent is therefore reused as the input to the next transition\. Small prediction errors introduced at one step can alter subsequent inputs and propagate throughout the rollout\. This is particularly undesirable for planning, where the terminal prediction may depend on several successive imagined transitions\.

Second, the predictor is trained using deterministic point\-wise alignment between a predicted latent and a clean target embedding\. This objective provides accurate supervision at the observed training samples, but does not explicitly constrain the behavior of the dynamics model under uncertainty in the latent representation\.\([Huang, 2026](https://arxiv.org/html/2608.29029#bib.bib28)\)If a visual perturbation shifts the encoded observation fromztz\_\{t\}tozt\+δz\_\{t\}\+\\delta, the predictor has not been directly trained to recover the same underlying transition structure from such a perturbed latent\.

## 4Method

### 4\.1Training Objective

Motivated by the limitations of deterministic autoregressive prediction, we formulate multi\-step latent dynamics as conditional flow matching\([Lipman et al\., 2022](https://arxiv.org/html/2608.29029#bib.bib9)\)\. Flow matching learns a continuous\-time vector field that transports samples from a source distributionp0p\_\{0\}toward a target distributionp1p\_\{1\}\. Given a source samplex0∼𝒩⁡\(0,I\)x\_\{0\}\\sim\\mathcal\{N\}\(0,I\), a target samplex1∼pdatax\_\{1\}\\sim p\_\{\\mathrm\{data\}\}, and flow timeτ∼𝒰⁡\[0,1\]\\tau\\sim\\mathcal\{U\}\[0,1\], a commonly used linear probability path is

xτ=\(1−τ\)​x0\+τ​x1,x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau x\_\{1\},\(7\)whose target velocity is constant along the path,

uτ​\(xτ∣x0,x1\)=x1−x0\.u\_\{\\tau\}\(x\_\{\\tau\}\\mid x\_\{0\},x\_\{1\}\)=x\_\{1\}\-x\_\{0\}\.\(8\)The corresponding conditional flow matching objective is

ℒCFM=𝔼x0∼𝒩⁡\(0,I\),x1∼pdata,τ∼𝒰⁡\[0,1\]​\[‖vθ​\(xτ,τ\)−\(x1−x0\)‖2\]\.\\mathcal\{L\}\_\{\\mathrm\{CFM\}\}=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}x\_\{0\}\\sim\\mathcal\{N\}\(0,I\),\\\\ x\_\{1\}\\sim p\_\{\\mathrm\{data\}\},\\\\ \\tau\\sim\\mathcal\{U\}\[0,1\]\\end\{subarray\}\}\\left\[\\left\\\|v\_\{\\theta\}\(x\_\{\\tau\},\\tau\)\-\(x\_\{1\}\-x\_\{0\}\)\\right\\\|^\{2\}\\right\]\.\(9\)Training therefore requires only regression of the vector field at randomly sampled intermediate states and does not require numerical ordinary differential equation \(ODE\) integration\. Importantly, an entire future sequence can be treated as a single joint variable, allowing all prediction horizons to evolve jointly rather than recursively feeding intermediate predictions into subsequent transitions\. In addition, training along stochastic source\-target paths exposes the predictor to perturbed intermediate states rather than only clean representations\.

Figure 2:Overview of F\-JEPA\.\(a\)training pipeline\. The conditional flow matching predictor transports 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\. Future latent tokens are processed jointly by Transformer blocks with flow\-time conditioning and are conditioned on the current latent state and temporally aligned action embeddings\.We instantiate this formulation for action\-conditioned latent world modeling, as illustrated in Figure[2](https://arxiv.org/html/2608.29029#S4.F2)\(a\)\. For a prediction horizon ofhh, each training sample containsh\+1h\+1observationsot,ot\+1,…,ot\+ho\_\{t\},o\_\{t\+1\},\\ldots,o\_\{t\+h\}andhhtemporally aligned action blocksat,at\+1,…,at\+h−1a\_\{t\},a\_\{t\+1\},\\ldots,a\_\{t\+h\-1\}\. Each action block contains five consecutive low\-level environment actions, matching the frame skip used in the dataset\. A shared visual encoderfϕf\_\{\\phi\}maps observations todd\-dimensional latent states,

zt=fϕ​\(ot\)∈ℝd,z\_\{t\}=f\_\{\\phi\}\(o\_\{t\}\)\\in\\mathbb\{R\}^\{d\},\(10\)while an action encoderhψh\_\{\\psi\}maps each action block to the same embedding dimension,

eta=hψ​\(at\)∈ℝd\.e\_\{t\}^\{a\}=h\_\{\\psi\}\(a\_\{t\}\)\\in\\mathbb\{R\}^\{d\}\.\(11\)The future latent trajectory and encoded action sequence are respectively defined as

Z=\[zt\+1,zt\+2,…,zt\+h\]∈ℝh×d,Z=\\left\[z\_\{t\+1\},z\_\{t\+2\},\\ldots,z\_\{t\+h\}\\right\]\\in\\mathbb\{R\}^\{h\\times d\},\(12\)Ea=\[eta,et\+1a,…,et\+h−1a\]∈ℝh×d\.E^\{a\}=\\left\[e\_\{t\}^\{a\},e\_\{t\+1\}^\{a\},\\ldots,e\_\{t\+h\-1\}^\{a\}\\right\]\\in\\mathbb\{R\}^\{h\\times d\}\.\(13\)
For each future position, we independently sample a source latent from a Gaussian distribution,

zt\+i0∼𝒩\(μ,σ2I\),i=1,…,h,z^\{0\}\_\{t\+i\}\\sim\\mathcal\{N\}\(\\mu,\\sigma^\{2\}I\),\\qquad i=1,\\ldots,h,\(14\)and stack the samples into

Z0=\[zt\+10,zt\+20,…,zt\+h0\]∈ℝh×d\.Z^\{0\}=\\left\[z^\{0\}\_\{t\+1\},z^\{0\}\_\{t\+2\},\\ldots,z^\{0\}\_\{t\+h\}\\right\]\\in\\mathbb\{R\}^\{h\\times d\}\.\(15\)
Conditioned on the current latent stateztz\_\{t\}, action sequenceEaE^\{a\}, and flow timeτ\\tau, the predictorvθ​\(Xτ,τ,zt,Ea\)v\_\{\\theta\}\(X\_\{\\tau\},\\tau,z\_\{t\},E^\{a\}\)estimates the velocity of the complete future latent trajectory\. Following the standard CFM formulation, the F\-JEPA training objective is

ℒFM=𝔼\(zt,Ea,Z\)∼pdata,Z0∼𝒩⁡\(μ,σ2​I\),τ∼𝒰\[0,1\)​\[‖vθ​\(Xτ,τ,zt,Ea\)−\(Z−Z0\)‖2\],\\mathcal\{L\}\_\{\\mathrm\{FM\}\}=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\(z\_\{t\},E^\{a\},Z\)\\sim p\_\{\\mathrm\{data\}\},\\\\ Z^\{0\}\\sim\\mathcal\{N\}\(\\mu,\\sigma^\{2\}I\),\\\\ \\tau\\sim\\mathcal\{U\}\[0,1\)\\end\{subarray\}\}\\left\[\\left\\\|v\_\{\\theta\}\(X\_\{\\tau\},\\tau,z\_\{t\},E^\{a\}\)\-\(Z\-Z^\{0\}\)\\right\\\|^\{2\}\\right\],\(16\)whereXτ=\(1−τ\)​Z0\+τ​ZX\_\{\\tau\}=\(1\-\\tau\)Z^\{0\}\+\\tau Z\.

Finally, we retain SIGReg from LeWM to prevent representation collapse\. The regularizer is applied to the embeddings of both the current and future observations in each sampled sequence\. The complete optimization objective is

ℒ=ℒFM\+λSIG​ℒSIG\.\\mathcal\{L\}=\\mathcal\{L\}\_\{\\mathrm\{FM\}\}\+\\lambda\_\{\\mathrm\{SIG\}\}\\mathcal\{L\}\_\{\\mathrm\{SIG\}\}\.\(17\)SIGReg constrains the global geometry of the learned representation space, while the flow objective learns action\-conditioned dynamics along stochastic source–target interpolation paths\.

At inference time, generation starts by drawing a fresh sample from the Gaussian distribution as a source trajectoryZ0∼𝒩⁡\(μ,σ2​I\)Z^\{0\}\\sim\\mathcal\{N\}\(\\mu,\\sigma^\{2\}I\)\. Conditioned onztz\_\{t\}and the action sequenceEaE^\{a\}, the learned vector field defines the ODEd​Xτ/d​τ=vθ​\(Xτ,τ,zt,Ea\)\\mathrm\{d\}X\_\{\\tau\}/\\mathrm\{d\}\\tau=v\_\{\\theta\}\(X\_\{\\tau\},\\tau,z\_\{t\},E^\{a\}\)\. We solve this ODE using explicit Euler integration\. GivenNNflow steps, the integration interval isΔ​τ=1/N\\Delta\\tau=1/N, and the latent trajectory is updated asXτ\+Δ​τ=Xτ\+Δ​τ​vθ​\(Xτ,τ,zt,Ea\)X\_\{\\tau\+\\Delta\\tau\}=X\_\{\\tau\}\+\\Delta\\tau\\,v\_\{\\theta\}\(X\_\{\\tau\},\\tau,z\_\{t\},E^\{a\}\)\. Starting fromX0=Z0X\_\{0\}=Z^\{0\}, the final integrated sampleX1X\_\{1\}is taken as the predicted future latent trajectoryZ^\\hat\{Z\}\. All future latent states are updated jointly at each integration step, rather than generated autoregressively\.

### 4\.2Model Architecture

Following LeWM, each observationoto\_\{t\}is encoded by a Vision Transformer \(ViT\)\([Dosovitskiy et al\., 2020](https://arxiv.org/html/2608.29029#bib.bib20)\), and its\[CLS\]representation is projected through a multilayer perceptron \(MLP\) to obtain the latent stateztz\_\{t\}\. Action blocks are independently mapped to action embeddings by the action encoder\. Figure[2](https://arxiv.org/html/2608.29029#S4.F2)\(b\) illustrates the conditional flow matching predictor used by F\-JEPA\. The dynamics predictor is a Transformer\-based\([Vaswani et al\., 2017](https://arxiv.org/html/2608.29029#bib.bib22)\)conditional flow matching predictor\. Each future latent token is augmented with a temporal position embedding identifying its prediction horizon\. The scalar flow timeτ\\tauis first mapped through a sinusoidal embedding and then injected into each Transformer block using adaptive layer normalization \(AdaLN\)\([Peebles and Xie, 2023](https://arxiv.org/html/2608.29029#bib.bib21)\)\. The future trajectory is conditioned on both the current observation latent and the encoded action sequence\. The current latent state is concatenated with the encoded action sequence to form the context for cross\-attention\. Additional implementation details, including the Gaussian source configuration and hyperparameters, are provided in Appendix[A](https://arxiv.org/html/2608.29029#A1)\.

## 5Experiments

### 5\.1Experimental Setup

#### Planning configurations\.

Following LeWM, we evaluate F\-JEPA on four environments: Two\-Room\([Sobal et al\., 2025](https://arxiv.org/html/2608.29029#bib.bib32)\), Reacher\([Tassa et al\., 2018](https://arxiv.org/html/2608.29029#bib.bib33)\), Push\-T\([Zhou et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib6)\), and OGBench\-Cube\([Park et al\., 2025](https://arxiv.org/html/2608.29029#bib.bib34)\)\. See Appendix[B\.1](https://arxiv.org/html/2608.29029#A2.SS1)for details of the environments and the datasets\. We use model\-predictive control \(MPC\)\([Richalet et al\., 1978](https://arxiv.org/html/2608.29029#bib.bib36);[Hansen et al\., 2022](https://arxiv.org/html/2608.29029#bib.bib35);[Hansen et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib37)\)with the cross\-entropy method \(CEM\)\([Rubinstein and Kroese, 2004](https://arxiv.org/html/2608.29029#bib.bib38)\)for planning\. At each planning step, CEM optimizes a sequence of 5 action blocks, each containing 5 low\-level actions\. Each refinement iteration evaluates 300 randomly sampled candidate sequences and updates the action distribution with the best 30\. Only the terminal predicted latent state is used to compute the planning cost\. The selected five\-block sequence is executed before replanning from the new observation\. We report success rates over 50 evaluation rollouts\. Planning settings follow LeWM except on Push\-T, where F\-JEPA uses 10 CEM iterations while the reported LeWM baseline uses its original 30\-iteration setting\. We examine the effect of this difference explicitly in Section[5\.3](https://arxiv.org/html/2608.29029#S5.SS3)\.

#### Visual perturbations\.

To evaluate robustness independently of the underlying task dynamics, we introduce a spatially localized Gaussian perturbation directly in RGB space\. Letppdenote a pixel location,kkan RGB channel, andccthe sampled perturbation center\. We drawϵp,k∼𝒩⁡\(0,s2\)\\epsilon\_\{p,k\}\\sim\\mathcal\{N\}\(0,s^\{2\}\)and construct

o~p,k=clip⁡\(op,k\+exp⁡\(−∥p−c∥222​r2\)​ϵp,k,0,255\),\\widetilde\{o\}\_\{p,k\}=\\operatorname\{clip\}\\left\(o\_\{p,k\}\+\\exp\\\!\\left\(\-\\frac\{\\lVert p\-c\\rVert\_\{2\}^\{2\}\}\{2r^\{2\}\}\\right\)\\epsilon\_\{p,k\},0,255\\right\),\(18\)wheresscontrols the noise magnitude andrrcontrols its spatial decay\. The center is sampled from the background, away from the agent and task\-relevant objects, and remains fixed within each rollout\. In the main comparison, we uses=100s=100andr=35r=35for Two\-Room, Reacher, and OGBench\-Cube\. For Push\-T, we user=10r=10, sincer=35r=35reduces the LeWM baseline to near\-random performance and yields a less informative comparison\.

### 5\.2Planning Performance and Robustness

Tables[1](https://arxiv.org/html/2608.29029#S5.T1)and[2](https://arxiv.org/html/2608.29029#S5.T2)summarize planning performance under clean and perturbed observations, while Figure[1](https://arxiv.org/html/2608.29029#S1.F1)provides a visual comparison with LeWM\. Under clean observations, F\-JEPA matches or exceeds LeWM across all four environments and achieves the highest mean success rate, improving the average from 86% to 92%\. Thus, joint flow\-based trajectory prediction improves robustness without sacrificing performance under clean conditions\.

Table 1:Planning success rates \(%\) under clean visual conditions\.F\-JEPA achieves the highest mean success rate \(92%\) and the best performance on three of the four environments, demonstrating strong planning performance without visual perturbations\. Results for PLDM\([Sobal et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib18)\), DINO\-WM\([Zhou et al\., 2024](https://arxiv.org/html/2608.29029#bib.bib6)\), and LeWM are taken from\([Maes et al\., 2026](https://arxiv.org/html/2608.29029#bib.bib7)\)\. The best result in each column is shown in bold\.Table 2:Planning success rates \(%\) under noisy visual conditions\.F\-JEPA improves performance across all four environments, increasing the mean success rate from 67% to 86% and demonstrating substantially greater robustness to visual perturbations\. The final row reports the absolute percentage\-point improvement over LeWM\.![Refer to caption](https://arxiv.org/html/2608.29029v1/tworoom_rollouts.png)Figure 3: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 while F\-JEPA remains successful, illustrating improved robustness to visual perturbations\.![Refer to caption](https://arxiv.org/html/2608.29029v1/cube_rollouts.png)Figure 4: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 while F\-JEPA remains successful, illustrating improved robustness to visual perturbations\.The advantage is substantially larger under visual perturbations\. F\-JEPA improves success in every environment, raising the mean from 67% to 86%, an average improvement of 19 percentage points\. Accordingly, the mean degradation from clean to noisy observations is reduced from 19 points for LeWM to 6 points for F\-JEPA\. These results are consistent with our hypothesis that stochastic flow training over perturbed latent trajectories promotes more robust latent dynamics\.

The qualitative rollouts in Figures[3](https://arxiv.org/html/2608.29029#S5.F3)and[4](https://arxiv.org/html/2608.29029#S5.F4)visualize the robustness\. Under clean observations, both LeWM and F\-JEPA successfully complete the tasks\. After task\-irrelevant visual perturbations are introduced, however, LeWM fails in both examples, whereas F\-JEPA remains successful\.

### 5\.3Ablation Studies

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 ats=100s=100, andΔ\\Deltadenotes the absolute percentage\-point improvement over LeWM\.Table 4:Reacher success rates \(%\) across Gaussian\-noise standard deviations\.F\-JEPA consistently outperforms LeWM across all tested noise levels and retains a larger performance margin under severe perturbations\. The patch radius is fixed atr=35r=35, andΔ\\Deltadenotes the absolute percentage\-point improvement over LeWM\.#### Robustness across perturbation severity\.

We first vary the spatial extent and magnitude of the visual perturbation on Reacher while holding all other settings fixed\. Table[3](https://arxiv.org/html/2608.29029#S5.T3)varies the radius withs=100s=100, and Table[4](https://arxiv.org/html/2608.29029#S5.T4)varies the noise standard deviation withr=35r=35\. F\-JEPA consistently outperforms LeWM across every tested configuration, showing that the robustness gain is not specific to a single perturbation setting\.

Table 5:Push\-T success rates \(%\) across flow integration steps and corresponding inference cost\.8 Euler steps achieve the highest noisy success rate while retaining the best clean performance, whereas additional integration steps substantially increase runtime without improving planning performance\. Times report wall\-clock seconds over 50 rollouts; bold denotes the highest success rate within each condition\.
#### Flow integration steps\.

Table[5](https://arxiv.org/html/2608.29029#S5.T5)studies the accuracy–computation trade\-off of ODE integration on Push\-T\. 8 Euler steps preserve the best clean performance while achieving the highest noisy success rate\. Increasing the number of integration steps beyond 8 substantially increases runtime without improving control performance, indicating that finer numerical integration is not the limiting factor in this setting\. We therefore use 8 flow steps in the main experiments\.

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\.
#### CEM refinement iterations\.

Finally, we examine whether the Push\-T improvement can be attributed to differences in CEM planning compute\. The main F\-JEPA result uses 10 CEM iterations, whereas the reported LeWM baseline retains its original 30\-iteration setting\. Table[6](https://arxiv.org/html/2608.29029#S5.T6)evaluates both methods under matched CEM budgets\. F\-JEPA maintains 96% clean success across the full sweep and achieves its highest noisy success with only 10 iterations\.

Additional ablation results and visualizations are provided in Appendix[B\.2](https://arxiv.org/html/2608.29029#A2.SS2)\.

## 6Conclusion

We presented F\-JEPA, a conditional flow matching dynamics model for reconstruction\-free JEPA world models\. F\-JEPA replaces deterministic one\-step autoregressive prediction with stochastic transport over the complete future latent trajectory, allowing all prediction horizons to evolve jointly\. Across four control environments, F\-JEPA improves average planning performance under clean observations and substantially reduces performance degradation under task\-irrelevant visual perturbations\. Ablations further show that the robustness improvement persists across perturbation strengths, does not require increasingly fine ODE integration or a larger CEM planning budget\. Together, these results support conditional flow matching as a promising alternative to deterministic autoregressive latent dynamics for JEPA\-based world modeling\.

Limitations and future work\.Our experiments focus on relatively short\-horizon control tasks and robustness is evaluated primarily under localized Gaussian visual perturbations\. Future work could investigate longer\-horizon tasks and broader distribution shifts\. Finally, flow\-based prediction requires multiple vector\-field evolutions at inference time\. Although our ablations show that a small number of integration steps is sufficient in the studied tasks, improving the efficiency of flow\-based latent prediction remains an important direction for future work\.

### AI use statement

In this work, we used generative AI tools to assist with formulating mathematical claims, providing feedback on experiments, implementing methods, and interpreting results\. We did not use generative AI tools to develop theoretical models or conceptual frameworks or to propose or refine hypotheses\. The following uses were not applicable to this work: generating synthetic datasets, providing critical ingredients for proving mathematical claims, assisting in the writing of proofs, assisting with translation, cleaning or reformatting datasets, and supporting qualitative or thematic data analysis\.

Additionally, we used generative AI tools to edit the research paper for improved readability, identify relevant literature, and propose a title for the paper\. We reviewed all AI\-assisted work\. LLM\-generated code was verified and tested for correctness by the authors, and LLM\-generated mathematical claims were examined and refined to ensure that they were correct and aligned with the actual method\. We take responsibility for the final content of this work, including all text, claims, and artifacts produced with the aid of generative AI\.

### Reproducibility statement

For reproducibility, Section[4](https://arxiv.org/html/2608.29029#S4)provides the complete training objective and model architecture\. Appendix[A](https://arxiv.org/html/2608.29029#A1)documents implementation details, training hyperparameters, Gaussian source configurations, and hardware used for all experiments\. The planning and evaluation protocols are specified in Section[5\.1](https://arxiv.org/html/2608.29029#S5.SS1), test environments and datasets are detailed in Appendix[B\.1](https://arxiv.org/html/2608.29029#A2.SS1)\. We also ablate key inference\-time parameters, including flow integration steps and CEM refinement iterations, to make the reported performance and computational settings directly reproducible\.

## References

- Alonsoet al\.\(2024\)E\. Alonso, A\. Jelley, V\. Micheli, A\. Kanervisto, A\. Storkey, T\. Pearce, and F\. FleuretDiffusion for world modeling: visual details matter in atari\.Advances in Neural Information Processing Systems37,pp\. 58757–58791\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Assranet al\.\(2023\)M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, Y\. LeCun, and N\. BallasSelf\-supervised learning from images with a joint\-embedding predictive architecture\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 15619–15629\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px2.p1.1)\.
- Assranet al\.\(2025\)M\. Assran, A\. Bardes, D\. Fan, Q\. Garrido, R\. Howes, M\. Muckley, A\. Rizvi, C\. Roberts, K\. Sinha, A\. Zholus,et al\.V\-jepa 2: self\-supervised video models enable understanding, prediction and planning\.arXiv preprint arXiv:2506\.09985\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px2.p1.1)\.
- Balestriero and LeCun \(2025\)R\. Balestriero and Y\. LeCunLejepa: provable and scalable self\-supervised learning without the heuristics\.arXiv preprint arXiv:2511\.08544\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p2.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px2.p1.1),[§3\.1](https://arxiv.org/html/2608.29029#S3.SS1.p3.1)\.
- Bardeset al\.\(2024\)A\. Bardes, Q\. Garrido, J\. Ponce, X\. Chen, M\. Rabbat, Y\. LeCun, M\. Assran, and N\. BallasRevisiting feature prediction for learning visual representations from video\.arXiv preprint arXiv:2404\.08471\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px2.p1.1)\.
- Blacket al\.\(2024\)K\. Black, N\. Brown, D\. Driess, A\. Esmail, M\. Equi, C\. Finn, N\. Fusai, L\. Groom, K\. Hausman, B\. Ichter,et al\.π0\\pi\_\{0\}: A vision\-language\-action flow model for general robot control\.arXiv preprint arXiv:2410\.24164\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p5.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px3.p1.1)\.
- Dosovitskiyet al\.\(2020\)A\. Dosovitskiy, L\. Beyer, A\. Kolesnikov, D\. Weissenborn, X\. Zhai, T\. Unterthiner, M\. Dehghani, M\. Minderer, G\. Heigold, S\. Gelly,et al\.An image is worth 16x16 words: transformers for image recognition at scale\.arXiv preprint arXiv:2010\.11929\.Cited by:[§4\.2](https://arxiv.org/html/2608.29029#S4.SS2.p1.1)\.
- Epps and Pulley \(1983\)T\. W\. Epps and L\. B\. PulleyA test for normality based on the empirical characteristic function\.Biometrika70\(3\),pp\. 723–726\.Cited by:[§3\.1](https://arxiv.org/html/2608.29029#S3.SS1.p3.2)\.
- Ha and Schmidhuber \(2018\)D\. Ha and J\. SchmidhuberWorld models\.arXiv preprint arXiv:1803\.101222\(3\),pp\. 440\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p1.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Hafneret al\.\(2019a\)D\. Hafner, T\. Lillicrap, J\. Ba, and M\. NorouziDream to control: learning behaviors by latent imagination\.arXiv preprint arXiv:1912\.01603\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p1.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Hafneret al\.\(2019b\)D\. Hafner, T\. Lillicrap, I\. Fischer, R\. Villegas, D\. Ha, H\. Lee, and J\. DavidsonLearning latent dynamics for planning from pixels\.InInternational conference on machine learning,pp\. 2555–2565\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Hafneret al\.\(2020\)D\. Hafner, T\. Lillicrap, M\. Norouzi, and J\. BaMastering atari with discrete world models\.arXiv preprint arXiv:2010\.02193\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Hafneret al\.\(2023\)D\. Hafner, J\. Pasukonis, J\. Ba, and T\. LillicrapMastering diverse domains through world models\.arXiv preprint arXiv:2301\.04104\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Hafneret al\.\(2025\)D\. Hafner, W\. Yan, and T\. LillicrapTraining agents inside of scalable world models\.arXiv preprint arXiv:2509\.24527\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Hansenet al\.\(2024\)N\. Hansen, H\. Su, and X\. WangTd\-mpc2: scalable, robust world models for continuous control\.InInternational Conference on Learning Representations,Vol\.2024,pp\. 47376–47405\.Cited by:[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1)\.
- Hansenet al\.\(2022\)N\. Hansen, X\. Wang, and H\. SuTemporal difference learning for model predictive control\.arXiv preprint arXiv:2203\.04955\.Cited by:[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1)\.
- Huang \(2026\)Y\. HuangVjepa: variational joint embedding predictive architectures as probabilistic world models\.arXiv preprint arXiv:2601\.14354\.Cited by:[§3\.2](https://arxiv.org/html/2608.29029#S3.SS2.p3.1)\.
- Jiaet al\.\(2026\)J\. Jia, G\. Li, X\. Chen, T\. An, Y\. Hu, J\. Li, X\. Guo, and J\. YangAction\-to\-action flow matching\.arXiv preprint arXiv:2602\.07322\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px3.p1.1)\.
- Kimet al\.\(2026\)J\. Kim, J\. Park, S\. Jeon, and S\. KimBetter source, better flow: learning condition\-dependent source distribution for flow matching\.arXiv preprint arXiv:2602\.05951\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px3.p1.1)\.
- LeCun \(2022\)Y\. LeCunA path towards autonomous machine intelligence version 0\.9\.2, 2022\-06\-27\.Open Review\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p2.1)\.
- Lipmanet al\.\(2022\)Y\. Lipman, R\. T\. Chen, H\. Ben\-Hamu, M\. Nickel, and M\. LeFlow matching for generative modeling\.arXiv preprint arXiv:2210\.02747\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p5.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px3.p1.1),[§4\.1](https://arxiv.org/html/2608.29029#S4.SS1.p1.1)\.
- Liu \(2022\)Q\. LiuRectified flow: a marginal preserving approach to optimal transport\.arXiv preprint arXiv:2209\.14577\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px3.p1.1)\.
- Liuet al\.\(2022\)X\. Liu, C\. Gong, and Q\. LiuFlow straight and fast: learning to generate and transfer data with rectified flow\.arXiv preprint arXiv:2209\.03003\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px3.p1.1)\.
- Maeset al\.\(2026\)L\. Maes, Q\. L\. Lidec, D\. Scieur, Y\. LeCun, and R\. BalestrieroLeworldmodel: stable end\-to\-end joint\-embedding predictive architecture from pixels\.arXiv preprint arXiv:2603\.19312\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p2.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px2.p1.1),[§3\.1](https://arxiv.org/html/2608.29029#S3.SS1.p3.1),[Table 1](https://arxiv.org/html/2608.29029#S5.T1)\.
- Micheliet al\.\(2022\)V\. Micheli, E\. Alonso, and F\. FleuretTransformers are sample\-efficient world models\.arXiv preprint arXiv:2209\.00588\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p1.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Parket al\.\(2025\)S\. Park, K\. Frans, B\. Eysenbach, and S\. LevineOgbench: benchmarking offline goal\-conditioned rl\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 94937–94982\.Cited by:[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1)\.
- Peebles and Xie \(2023\)W\. Peebles and S\. XieScalable diffusion models with transformers\.In2023 IEEE/CVF International Conference on Computer Vision \(ICCV\),pp\. 4172–4182\.Cited by:[§4\.2](https://arxiv.org/html/2608.29029#S4.SS2.p1.1)\.
- Porcheret al\.\(2026\)F\. Porcher, N\. Carion, K\. Alahari, and S\. ChenFlow matching in feature space for stochastic world modeling\.arXiv preprint arXiv:2606\.29059\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Richaletet al\.\(1978\)J\. Richalet, A\. Rault, J\. Testud, and J\. PaponModel predictive heuristic control: applications to industrial processes\.Automatica14\(5\),pp\. 413–428\.Cited by:[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1)\.
- Rubinstein and Kroese \(2004\)R\. Y\. Rubinstein and D\. P\. KroeseThe cross\-entropy method: a unified approach to combinatorial optimization, monte\-carlo simulation and machine learning\.Springer New York\.Cited by:[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1)\.
- Sobalet al\.\(2026\)U\. Sobal, W\. Zhang, K\. Cho, R\. Balestriero, T\. G\. Rudner, and Y\. LeCunLearning from reward\-free offline data: a case for planning with latent dynamics models\.Advances in Neural Information Processing Systems38,pp\. 43905–43941\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px2.p1.1),[Table 1](https://arxiv.org/html/2608.29029#S5.T1)\.
- Sobalet al\.\(2025\)V\. Sobal, W\. Zhang, K\. Cho, R\. Balestriero, T\. G\. J\. Rudner, and Y\. LeCunStress\-testing offline reward\-free reinforcement learning: a case for planning with latent dynamics models\.InICLR 2025 Workshop on World Models: Understanding, Modelling and Scaling,Cited by:[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1)\.
- Sunet al\.\(2026\)J\. Sun, W\. Zhang, Z\. Qi, S\. Ren, Z\. Liu, H\. Zhu, G\. Sun, X\. Jin, and Z\. ChenVla\-jepa: enhancing vision\-language\-action model with latent world model\.arXiv preprint arXiv:2602\.10098\.Cited by:[§3\.2](https://arxiv.org/html/2608.29029#S3.SS2.p1.1)\.
- Tassaet al\.\(2018\)Y\. Tassa, Y\. Doron, A\. Muldal, T\. Erez, Y\. Li, D\. d\. L\. Casas, D\. Budden, A\. Abdolmaleki, J\. Merel, A\. Lefrancq,et al\.Deepmind control suite\.arXiv preprint arXiv:1801\.00690\.Cited by:[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1)\.
- Vaswaniet al\.\(2017\)A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, Ł\. Kaiser, and I\. PolosukhinAttention is all you need\.Advances in neural information processing systems30\.Cited by:[§4\.2](https://arxiv.org/html/2608.29029#S4.SS2.p1.1)\.
- Walkeret al\.\(2025\)J\. C\. Walker, P\. Vélez, L\. P\. Cabrera, G\. Zhou, S\. Ebrahimi, R\. Kabra, C\. Doersch, M\. Ovsjanikov, J\. Carreira, and S\. GinosarFrozen forecasting: a unified evaluation\.arXiv preprint arXiv:2507\.13942\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px1.p1.1)\.
- Wanet al\.\(2025\)T\. Wan, A\. Wang, B\. Ai, B\. Wen, C\. Mao, C\. Xie, D\. Chen, F\. Yu, H\. Zhao, J\. Yang,et al\.Wan: open and advanced large\-scale video generative models\.arXiv preprint arXiv:2503\.20314\.Cited by:[§1](https://arxiv.org/html/2608.29029#S1.p5.1),[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px3.p1.1)\.
- Zhouet al\.\(2024\)G\. Zhou, H\. Pan, Y\. LeCun, and L\. PintoDino\-wm: world models on pre\-trained visual features enable zero\-shot planning\.arXiv preprint arXiv:2411\.04983\.Cited by:[§2](https://arxiv.org/html/2608.29029#S2.SS0.SSS0.Px2.p1.1),[§5\.1](https://arxiv.org/html/2608.29029#S5.SS1.SSS0.Px1.p1.1),[Table 1](https://arxiv.org/html/2608.29029#S5.T1)\.

## Appendix AImplementation Details

#### Flow source and attention configuration\.

We consider two Gaussian source parameterizations\. The*standard\-noise*source independently samples each future token aszt\+i0∼𝒩⁡\(0,I\)z^\{0\}\_\{t\+i\}\\sim\\mathcal\{N\}\(0,I\), whereas the*state\-centered*source sampleszt\+i0∼𝒩⁡\(zt,0\.52​I\)z^\{0\}\_\{t\+i\}\\sim\\mathcal\{N\}\(z\_\{t\},0\.5^\{2\}I\)\. We further consider either bidirectional or causal self\-attention among future trajectory tokens\. Under causal self\-attention, the token at future positioniiattends only to itself and preceding future positions, whereas bidirectional self\-attention allows interactions across the complete predicted trajectory\. In both variants, action conditioning remains temporally causal: the token at positioniireceives only the action prefixat:t\+i−1a\_\{t:t\+i\-1\}\. Two\-Room and OGBench\-Cube use the standard\-noise source with bidirectional future\-token self\-attention, while Reacher and Push\-T use the state\-centered source with causal future\-token self\-attention\.

#### Architecture\.

All observations are resized to224×224224\\times 224and normalized before being passed to the visual encoder\. We use a ViT\-Tiny with14×1414\\times 14patches as the visual encoder and an MLP as the action encoder\. Both projected visual latents and action embeddings have dimension 192\. For the flow matching predictor, we use 6 Transformer blocks with 16 attention heads, head dimension 64, and an MLP hidden dimension of 2048\.

#### Training details\.

All F\-JEPA models are trained for 20 epochs with a batch size of 128, a prediction horizon of 5, and a history size of 1\. For Push\-T, the selected model initializes the visual encoder and projector from a trained LeWM checkpoint and keeps both components frozen during training; predictor dropout is set to 0\. For Reacher, the visual encoder and projector are also initialized from a LeWM checkpoint but remain trainable as a warm start\. The remaining selected models learn the visual representation jointly from training and use a predictor dropout of 0\.1\. The SIGReg weight is set to 0\.09 for Push\-T and 0\.1 for the remaining environments\. All other preprocessing, optimization, and training settings follow LeWM\.

#### Hardware\.

All training and evaluation experiments are conducted on a single NVIDIA RTX 6000 Ada GPU\.

## Appendix BDatasets and Additional Results

### B\.1Test Environments and Datasets

We evaluate on four visually distinct control environments covering navigation and manipulation, illustrated in Figure[5](https://arxiv.org/html/2608.29029#A2.F5): Two\-Room, a goal\-directed navigation task in a two\-room layout; Reacher, a continuous\-control task requiring a planar arm to reach a target; Push\-T, a planar manipulation task in which an agent pushes a T\-shaped object toward a goal configuration; and OGBench\-Cube, a goal\-conditioned cube manipulation task\.

![Refer to caption](https://arxiv.org/html/2608.29029v1/environment_first_frames.png)Figure 5: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\.Following the LeWM dataset construction, observations are recorded with a frame skip of 5, such that two consecutive observation frames are separated by 5 low\-level environment steps\. The corresponding low\-level actions are concatenated into an action block, containing the 5 actions executed between adjacent observations\. Consequently, each latent transition in the world model is conditioned on one action block rather than a single low\-level action\.

### B\.2Additional Results

![Refer to caption](https://arxiv.org/html/2608.29029v1/reacher_rollouts.png)Figure 6: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\.![Refer to caption](https://arxiv.org/html/2608.29029v1/pusht_rollouts.png)Figure 7: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\.#### Qualitative rollouts\.

Figures[6](https://arxiv.org/html/2608.29029#A2.F6)and[7](https://arxiv.org/html/2608.29029#A2.F7)provide additional comparisons between LeWM and F\-JEPA\. Consistent with the main\-text examples, F\-JEPA remains more reliable after visual perturbations are introduced\. In Reacher, both methods complete the task under the clean condition, but F\-JEPA reaches the target with higher precision\. Together, these rollouts support the improved planning performance and robustness observed across the four environments\.

Table 7:Additional Gaussian\-patch radius ablations\.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 ats=100s=100, andΔ\\Deltadenotes the absolute percentage\-point difference between F\-JEPA and LeWM\.Figure 8: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](https://arxiv.org/html/2608.29029#S5.T3)and[7](https://arxiv.org/html/2608.29029#A2.T7)\.
#### Robustness ablations\.

To complement the Reacher perturbation ablations reported in the main text, we evaluate the remaining environments across Gaussian\-patch radii and noise magnitudes\. These results test whether the robustness trends observed in the main experiments persist across different perturbation scales and environments\.

Table 8:Additional Gaussian\-noise magnitude ablations\.Across Two\-Room, Push\-T, and OGBench\-Cube, F\-JEPA generally maintains higher success than LeWM as the perturbation magnitude increases, demonstrating robustness beyond the noise level used in the main evaluation\. The patch radius is fixed atr=35r=35for Two\-Room and OGBench\-Cube and atr=10r=10for Push\-T, andΔ\\Deltadenotes the absolute percentage\-point difference between F\-JEPA and LeWM\.Figure 9: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[4](https://arxiv.org/html/2608.29029#S5.T4)and[8](https://arxiv.org/html/2608.29029#A2.T8)\.Tables[7](https://arxiv.org/html/2608.29029#A2.T7)and[8](https://arxiv.org/html/2608.29029#A2.T8)show that the robustness gains are not tied to a single perturbation configuration\. Across the radius sweep, F\-JEPA matches or exceeds LeWM at every tested setting in the three environments\. The noise\-magnitude sweep exhibits a similar overall trend: F\-JEPA remains stronger across nearly all settings, including the largest tested perturbations, with only a small reversal on OGBench\-Cube ats=150s=150\. Push\-T becomes particularly challenging as the perturbation grows spatially, yet F\-JEPA retains an advantage or matches the baseline throughout the sweep\. Together with the Reacher results in the main text, these experiments show that the observed robustness improvement persists across a broad range of visual perturbation strengths and spatial scales\. Figures[8](https://arxiv.org/html/2608.29029#A2.F8)and[9](https://arxiv.org/html/2608.29029#A2.F9)summarize the corresponding success\-rate trends across all four environments\.

#### Flow\-step ablations\.

We additionally evaluate the effect of ODE integration steps on Two\-Room, Reacher, and OGBench\-Cube, complementing the Push\-T ablation reported in the main text\.

Figure 10: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](https://arxiv.org/html/2608.29029#S5.T5)and[9](https://arxiv.org/html/2608.29029#A2.T9)\.Table 9:Additional flow integration\-step ablations\.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\.Table[9](https://arxiv.org/html/2608.29029#A2.T9)confirms the trend observed on Push\-T\. Increasing the number of Euler steps does not monotonically improve planning performance: Two\-Room is already near saturation with 4–8 steps, while both Reacher and OGBench\-Cube achieve their strongest clean and noisy results with 8 steps\. In contrast, inference time grows substantially as the number of integration steps increases\. These results support 8 Euler steps as a favorable accuracy\-efficiency trade\-off across the evaluated environments and indicate that finer numerical integration is not the primary bottleneck for planning performance\. Figure[10](https://arxiv.org/html/2608.29029#A2.F10)summarizes the effect of the number of Euler integration steps across all four environments\.

#### Inference\-time analysis\.

We compare the inference cost of F\-JEPA and LeWM on Two\-Room using the same 10\-iteration CEM planning budget\. F\-JEPA requires multiple vector\-field evaluations to integrate the flow trajectory, whereas LeWM autoregressively applies its one\-step predictor across the 5\-step prediction horizon\. Table[10](https://arxiv.org/html/2608.29029#A2.T10)reports both the latency of the dynamics model and the resulting end\-to\-end evaluation time\.

Table 10:Inference\-time comparison between LeWM and F\-JEPA\.F\-JEPA incurs higher dynamics and per\-planning latency due to iterative flow integration, but the end\-to\-end runtime gap over 50 rollouts is considerably smaller because LeWM requires substantially more replanning during task execution\. Both methods use 10 CEM refinement iterations\.A single flow step is more expensive than a one\-step LeWM prediction \(5\.0 ms versus 2\.1 ms\)\. Moreover, one F\-JEPA planning cycle performs 8 flow integration steps over 10 CEM iterations, resulting in 400\.0 ms per planning step, compared with 105\.0 ms for the 5 autoregressive predictions used by LeWM\. Thus, F\-JEPA is approximately3\.8×3\.8\\timesslower per planning cycle\. However, this difference does not translate directly to end\-to\-end evaluation time: over 50 rollouts, F\-JEPA takes 34\.4 s compared with 21\.7 s for LeWM, corresponding to only a1\.6×1\.6\\timesincrease\. This smaller practical gap arises because LeWM requires substantially more replanning steps over complete rollouts, whereas F\-JEPA more often completes the task with fewer replanning cycles\. These results highlight that the additional cost of iterative flow integration is partially offset at the task level by more effective planning\. Moreover, because F\-JEPA updates the complete future trajectory jointly at each flow step rather than recursively applying a one\-step predictor, its relative computational trade\-off may become more favorable for longer prediction horizons, where autoregressive rollout requires an increasing number of sequential model evaluations\.