Toward Goal-Agnostic Joint-Embedding Predictive Control of Partial Differential Equations
Summary
This paper presents a goal-agnostic control framework for partial differential equations using a joint-embedding predictive architecture (JEPA) with a lightweight 2D ViT encoder and action-conditioned latent dynamics, showing that using a learned physical observable probe outperforms raw latent distance for control tasks.
View Cached Full Text
Cached at: 07/27/26, 07:40 AM
# Toward Goal-Agnostic Joint-Embedding Predictive Control of Partial Differential Equations
Source: [https://arxiv.org/html/2607.21644](https://arxiv.org/html/2607.21644)
Roberto Guglielmi University of Waterloo roberto\.guglielmi@uwaterloo\.ca †Corresponding author
###### Abstract
We present a goal\-agnostic control framework for partial differential equations \(PDEs\) built around an end to end joint\-embedding predictive architecture \(JEPA\)\. A lightweight 2D ViT encoder and action\-conditioned latent dynamics are trained offline without a reward or downstream goal, frozen, and reused by a model\-predictive path integral \(MPPI\) controller\. We find that when available, the control objective is better applied to an explicit physical observable \(provided injectivity\) than to minimizing raw Euclidean distance \(L2L^\{2\}\) in the learned latent space\. For a learned linear kinetic\-energy \(KE\) probe on frozen latent rollouts we can reproduce the control of held\-out trajectories withR2=0\.989R^\{2\}=0\.989, while requiring no change to the underlying world model\. On the PDE Control Gym 2D Navier–Stokes benchmark, using KE\-probe planning improves the matched 50\-episode native reward from−12\.08±0\.86\-12\.08\\pm 0\.86for latent\-L2L^\{2\}planning to−10\.90±0\.91\-10\.90\\pm 0\.91\(95% CI\), while lowering last\-quarter velocity\-field RMSE from0\.07650\.0765to0\.06920\.0692\. Across three intentionally withheld, dissimilar, aperiodic targets, KE planning lowers late field RMSE by53%53\\%relative to latent\-L2L^\{2\}planning \(0\.02200\.0220versus0\.04690\.0469\), winning all 30 paired episodes\. The same frozen model also supports controls targeting stabilization around a steady configuration via direct regulation of KE achieving2\.7%2\.7\\%mean relative error\. While the latent probe is brittle to measurement noise and missing pixels, we believe the results support the claim that latent dynamics can remain both dynamic and goal\-agnostic while calibrated observables \(granted they guarantee unique continuation\) may be a better objective for state control\.
## 1Introduction
Partial differential equations \(PDE\) arise in many applications, ranging from climate modeling, biological processes, to chemical process monitoring\. Developing control strategies for PDE\-systems can be challenging, as the states are high dimensional, and the relationships between quantities often nonlinear\. As such, the evaluation of candidate control\-sequences often requires computationally costly numerical simulation, typically precluding online planning and control\. Extant methods for PDE control, such as model based control, learned surrogate control, or reinforcement learning approaches, have been applied with varying degrees of success\. Classical model based controllers require problem specific analysis, or otherwise clever construction of reduced order models \(ROMs\)\[[6](https://arxiv.org/html/2607.21644#bib.bib43)\]\. Learned surrogate approaches introduce neural networks as approximations of the evolution operator for system dynamics\[[18](https://arxiv.org/html/2607.21644#bib.bib25),[19](https://arxiv.org/html/2607.21644#bib.bib26),[26](https://arxiv.org/html/2607.21644#bib.bib27)\], but quite often the surrogate is trained on the task of state space forward\-simulation and optimized against a specific objective, inheriting the challenges of full state reconstruction and requiring retraining when objectives change\[[35](https://arxiv.org/html/2607.21644#bib.bib28),[37](https://arxiv.org/html/2607.21644#bib.bib6)\]\. Reinforcement learning has demonstrated excellent performance on well defined tasks\[[7](https://arxiv.org/html/2607.21644#bib.bib40),[10](https://arxiv.org/html/2607.21644#bib.bib41),[28](https://arxiv.org/html/2607.21644#bib.bib38),[32](https://arxiv.org/html/2607.21644#bib.bib39)\], but suffers from the need to carefully select the reward function, which limits the ability to generalize\.
Joint\-embedding predictive architectures \(JEPAs\) instead focus on learning how applied actions map between elements of a learned latent space or "World Model"\. This is done without necessitating state reconstruction of any kind\[[2](https://arxiv.org/html/2607.21644#bib.bib2)\]\. Recent work in latent world modeling has used such reconstruction\-free features for zero\-shot planning primarily in visual domains\[[39](https://arxiv.org/html/2607.21644#bib.bib5)\], while exisiting latent PDE work indicates that compressed dynamics may accelerate both simulation and inverse optimization\[[37](https://arxiv.org/html/2607.21644#bib.bib6)\]\. We believe JEPA\-style latent world models offer unique benefits for PDE problems\. Training with random masking can potentially enable an encoder to consume both full and incomplete fields while still mapping to the same latent representation\. All the while, an action\-conditioned latent predictor enables goal\-independent model\-predictive control\. The learned dynamics model remains agnostic to the run\-time goal, since changing this objective at the time of deployment does not require retraining either the encoder or predictor\.
Latent planning is attractive for many reasons: its computational efficiency, and goal agnosticism being the two most obvious\. Unfortunately, a latent representation trained for predictability and action sensitivity is not automatically a calibrated metric space\. It is well known thatL2L^\{2\}distance between embeddings can distort both the scale and ordering of physical field error\. As such, we use temporal straightening to improve the latent rollout geometry\[[33](https://arxiv.org/html/2607.21644#bib.bib9)\], but our experiments show that regularization alone does not suffice to solve this issue entirely\. Instead, we fit a lightweight readout for a physical state observable and form the tracking error*after*that readout, in physical units\. Our observable is chosen carefully to guarantee unique continuation, and we observe that this controller substantially improves clean closed\-loop tracking without retraining or decoding the world model\.
Recent work using joint\-embedding world models has explored planning in visual manipulation, navigation, locomotion, and aerial domains\[[1](https://arxiv.org/html/2607.21644#bib.bib21),[39](https://arxiv.org/html/2607.21644#bib.bib5),[24](https://arxiv.org/html/2607.21644#bib.bib8),[38](https://arxiv.org/html/2607.21644#bib.bib7),[29](https://arxiv.org/html/2607.21644#bib.bib11)\]\. To our knowledge, a JEPA\-like joint\-embedding world model has yet to be applied to the closed\-loop control of PDE systems\. This work aims to demonstrate the potential for joint\-embedding architectures in PDE control, showing that a goal\-independent, joing\-embedding with a downstream predictor can be a uniquely dynamic controller, capable of natively supporting multiple PDE\-control objectives\.
We validate this framework in the PDE Control Gym two\-dimensional Navier–Stokes environment\[[4](https://arxiv.org/html/2607.21644#bib.bib1)\]\. We train masked joint\-embedding dynamics, enforce non\-collapse with VICReg\[[3](https://arxiv.org/html/2607.21644#bib.bib3)\], add temporal straightening and Delta\-JEPA\-like latent action decoding\[[38](https://arxiv.org/html/2607.21644#bib.bib7)\], and use MPPI with the frozen backbone\[[36](https://arxiv.org/html/2607.21644#bib.bib4)\]\.
Our Contributions:
1. 1\.To our knowledge, we introduce the first implementation of closed\-loop PDE\-control using a joint\-embedding latent world model, capable of learning action\-conditioned dynamics independently of the downstream control task\.
2. 2\.We introduce an observable\-aligned controller: a cheap to train linear probe estimates kinetic energy from predicted latents, and MPPI evaluates tracking error in that physical scalar rather than using raw latent distance\.
3. 3\.On the PDE Control Gym benchmark, this improves both native reward and velocity\-field RMSE over the controller using latent\-L2L^\{2\}cost\. Across three distinct aperiodic signals, it reduces late field RMSE by53%53\\%, with improvement in all 30 paired episodes\.
4. 4\.We reuse one frozen checkpoint and one five\-step latent horizon for time\-dependent tracking, fixed\-field stabilization, and kinetic energy control, ever working in state space or performing task\-specific fine tuning\.
## 2Goal\-Agnostic Control Formulation
### 2\.1Dynamics learned without a goal
Let a controlled field evolve asUt\+1=F\(Ut,at\)U\_\{t\+1\}=F\(U\_\{t\},a\_\{t\}\)\. We learn an encoderEθE\_\{\\theta\}and action\-conditioned latent predictorgϕg\_\{\\phi\}entirely from offline state transitions, without providing a reward, reference trajectory, equilibrium, or physical observable\. After training, the encoder and action\-conditioned predictor remain frozen, while the MPPI planner supplies sets of candidate action sequences, whose effect on the latent state are able to be probed throughgϕg\_\{\\phi\}\. This allows for a goal to be specified within the MPPI planner at the time of deployment, without touching the weights in eitherEθE\_\{\\theta\}, orgϕg\_\{\\phi\}\.
### 2\.2Deployment\-time goal interfaces
The flexibility of the planning machinery supports arbitrary control objectives\. We study three:
1. 1\.Control Gym Action Penalized Tracking\.A target sequence supplied as encoded stateszt\+1:t\+H⋆z^\{\\star\}\_\{t\+1:t\+H\}or physical observationsqt\+1:t\+H⋆q^\{\\star\}\_\{t\+1:t\+H\}can be tracked while making action\-penalty tradeoffs, achieving comparable performance to RL baselines without being trained in a cost\-aware manner\.
2. 2\.Time\-dependent tracking\.A target sequence can be supplied either as encoded stateszt\+1:t\+H⋆z^\{\\star\}\_\{t\+1:t\+H\}or as a time\-dependent physical observableqt\+1:t\+H⋆q^\{\\star\}\_\{t\+1:t\+H\}\.
3. 3\.Fixed\-target stabilization\.A specified state configuration is supplied and encoded only once\. The encoded latentz⋆z^\{\\star\}is treated as a constant target against which controls are evaluated, resulting in an equilibrium objective\.
Only the user\-specified cost changes between these tasks, meaning the action\-conditioned latent dynamics model and MPPI optimizer are otherwise unchanged\.
### 2\.3PDE Control Gym benchmark
We evaluate our model using the PDE Control Gym two\-dimensional Navier–Stokes environment at the configuration used by its published examples\[[4](https://arxiv.org/html/2607.21644#bib.bib1)\]\. The domain is a unit square, discretized withΔx=Δy=0\.05\\Delta x=\\Delta y=0\.05into a21×2121\\times 21grid\. The simulator usesΔt=10−3\\Delta t=10^\{\-3\}, horizonT=0\.2T=0\.2, viscosityν=0\.1\\nu=0\.1, density11, and 2000 pressure iterations\. The observation is the velocity fieldUt=\(ut,vt\)∈ℝ21×21×2U\_\{t\}=\(u\_\{t\},v\_\{t\}\)\\in\\mathbb\{R\}^\{21\\times 21\\times 2\}\. A scalar actionat∈\[0,4\]a\_\{t\}\\in\[0,4\]sets a controllable Dirichlet condition \(in the positive x direction\) along the upper boundary\.
The reference benchmark reward is
rt=−12HW∥Ut−Ut⋆∥22−γ2\(at−atref\)2,r\_\{t\}=\-\\frac\{1\}\{2HW\}\\lVert U\_\{t\}\-U\_\{t\}^\{\\star\}\\rVert\_\{2\}^\{2\}\-\\frac\{\\gamma\}\{2\}\(a\_\{t\}\-a\_\{t\}^\{\\mathrm\{ref\}\}\)^\{2\},\(1\)whereH=W=21H=W=21the resolution,γ=0\.1\\gamma=0\.1the action penalty, andatref=2a\_\{t\}^\{\\mathrm\{ref\}\}=2, the reference velocity against which actions are measured\. We report cumulative native reward, state\-space RMSE, and action total variation \(TV\)\. Our controller acts once per two native solver steps, so 99 macro actions cover 198 of the benchmark’s 199 scored native steps\.
### 2\.4Masked observations
When training our encoder \(lightweight ViT\), state observations are randomly masked\. Any state has probability0\.150\.15that it is fully observed; otherwise the masked fraction is sampled uniformly on the range\[0\.3,0\.9\]\[0\.3,0\.9\]\. The masks are selected from either contiguous rectangles, random pixels, or row strips with respective probabilities0\.50\.5,0\.30\.3, and0\.20\.2\. At evaluation, the current measurement is either fully observed, corrupted by zero\-mean Gaussian noise with standard deviation equal to55–15%15\\%of each velocity channel’s recent RMS, or missing55–25%25\\%of its pixels at random\. The goal field is always fully specified, where any target is interpreted as a specified objective, not a sensor measurement\.
An example of the masking protocol is shown in Figure[1](https://arxiv.org/html/2607.21644#S2.F1)\.
Figure 1:Masking protocol: contiguous rectangles, random pixels, and row strips, each shown hiding roughly 60% of the field\. Grey denotes masked \(zero\-filled\) pixels\.
## 3Goal\-Agnostic JEPA Control Framework
### 3\.1Masked latent dynamics
LetEθE\_\{\\theta\}denote a masked field encoder and letzt=Eθ\(Ut,mt\)z\_\{t\}=E\_\{\\theta\}\(U\_\{t\},m\_\{t\}\)\. The encoder is a compact two\-dimensional vision transformer with 256\-dimensional output,3×33\\times 3patches, four transformer blocks, and four attention heads\. The predictorgϕg\_\{\\phi\}consumes a history of four latent states and their actions, then auto\-regressively rolls out twelve latent steps at stride two\. The predictor consists of a two\-layer MLP with 512 hidden units per layer, trained to predict latent residuals rather than full states\.
For a given training sample the encoder is given two tasks\. First, it must encode the masked history, after which it must encode a second unmasked sequence containing the future fields\.
zt−K\+1:tctx=Eθ\(Ut−K\+1:t,m\),zt:t\+12real=Eθ\(Ut:t\+12,𝟏\)\.z^\{\\mathrm\{ctx\}\}\_\{t\-K\+1:t\}=E\_\{\\theta\}\(U\_\{t\-K\+1:t\},m\),\\qquad z^\{\\mathrm\{real\}\}\_\{t:t\+12\}=E\_\{\\theta\}\(U\_\{t:t\+12\},\\mathbf\{1\}\)\.The predictor is then tasked with rolling forward fromzctxz^\{\\mathrm\{ctx\}\}, where the unmaskedzrealz^\{\\mathrm\{real\}\}values are the latent targets\. In this setup, rollout gradients may flow through the predictor, the masked\-history encoder branch, and the unmasked\-future encoder branch\. The loss is thus
ℒroll=112∑h=112∥z^t\+h−zt\+hreal∥22\.\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{roll\}\}=\\frac\{1\}\{12\}\\sum\_\{h=1\}^\{12\}\\lVert\\hat\{z\}\_\{t\+h\}\-z^\{\\mathrm\{real\}\}\_\{t\+h\}\\rVert\_\{2\}^\{2\}\.\(2\)
We employ both VICReg variance and covariance penalties to discourage collapse\[[3](https://arxiv.org/html/2607.21644#bib.bib3)\], together with a temporal straightening loss that discourages latent trajectories from being unnecessarily circuitous\. An auxiliary decoderdψd\_\{\\psi\}is tasked with reconstructing the applied action from adjacent displacements in the unmasked, encoded real sequence \(not from predicted rollout displacements\),
ℒΔa=112∑h=112∥dψ\(zt\+hreal−zt\+h−1real\)−at\+h−1∥22,\\displaystyle\\mathcal\{L\}\_\{\\Delta a\}=\\frac\{1\}\{12\}\\sum\_\{h=1\}^\{12\}\\lVert d\_\{\\psi\}\(z^\{\\mathrm\{real\}\}\_\{t\+h\}\-z^\{\\mathrm\{real\}\}\_\{t\+h\-1\}\)\-a\_\{t\+h\-1\}\\rVert\_\{2\}^\{2\},\(3\)following the action\-sensitivity motivation of Delta\-JEPA\[[38](https://arxiv.org/html/2607.21644#bib.bib7)\]\. The temporal\-straightening penalty is likewise evaluated on the unmasked, encoded real sequence, not on the predicted rollout path\. The full composite loss is
ℒtotal=ℒroll\+λvarℒvar\+λcovℒcov\+λstrℒstr\+λΔaℒΔa\\mathcal\{L\}\_\{\\mathrm\{total\}\}=\\mathcal\{L\}\_\{\\mathrm\{roll\}\}\+\\lambda\_\{\\mathrm\{var\}\}\\,\\mathcal\{L\}\_\{\\mathrm\{var\}\}\+\\lambda\_\{\\mathrm\{cov\}\}\\,\\mathcal\{L\}\_\{\\mathrm\{cov\}\}\+\\lambda\_\{\\mathrm\{str\}\}\\,\\mathcal\{L\}\_\{\\mathrm\{str\}\}\+\\lambda\_\{\\Delta a\}\\,\\mathcal\{L\}\_\{\\Delta a\}\(4\)
with weightsλvar=1\.0\\lambda\_\{\\mathrm\{var\}\}=1\.0,λcov=1\.0\\lambda\_\{\\mathrm\{cov\}\}=1\.0,λstr=0\.1\\lambda\_\{\\mathrm\{str\}\}=0\.1, andλΔa=10\\lambda\_\{\\Delta a\}=10\. No decoder maps latent statezzback to the fieldUU\. The full training process is summarized in Figure[2](https://arxiv.org/html/2607.21644#S3.F2)\.
Figure 2:Joint\-embedding training process, in which ground\-truth and partially masked observations are encoded and rolled forward before their latents are compared\.
### 3\.2Latent MPPI
Given a candidate action sequenceA=\(a0,…,aH−1\)A=\(a\_\{0\},\\ldots,a\_\{H\-1\}\), the frozen predictor producesz^t\+1:t\+H\\hat\{z\}\_\{t\+1:t\+H\}\. MPPI samples 512 action sequences for three iterations and executes the first action of the exponentially weighted plan\[[36](https://arxiv.org/html/2607.21644#bib.bib4)\]\. Conventional latent\-space tracking uses
Jlat\(A\)=α∑h∈𝒯∥z^t\+h−zt\+h⋆∥22\+λa∑h=0H−1\(ah−2\)2,J\_\{\\mathrm\{lat\}\}\(A\)=\\alpha\\sum\_\{h\\in\\mathcal\{T\}\}\\lVert\\hat\{z\}\_\{t\+h\}\-z^\{\\star\}\_\{t\+h\}\\rVert\_\{2\}^\{2\}\+\\lambda\_\{a\}\\sum\_\{h=0\}^\{H\-1\}\(a\_\{h\}\-2\)^\{2\},\(5\)where𝒯\\mathcal\{T\}is the scored rollout tail\. Boundary effects propagate through the domain, so later predicted states more clearly distinguish candidate controls\. Our observable\-aligned alternative fits a scalar probeqωq\_\{\\omega\}on frozen predictor\-rollout latents and instead uses
Jq\(A\)=β∑h∈𝒯\[qω\(z^t\+h\)−qt\+h⋆\]2\+λa∑h=0H−1\(ah−2\)2\.J\_\{q\}\(A\)=\\beta\\sum\_\{h\\in\\mathcal\{T\}\}\\left\[q\_\{\\omega\}\(\\hat\{z\}\_\{t\+h\}\)\-q^\{\\star\}\_\{t\+h\}\\right\]^\{2\}\+\\lambda\_\{a\}\\sum\_\{h=0\}^\{H\-1\}\(a\_\{h\}\-2\)^\{2\}\.\(6\)For kinetic energy tracking,qω=k^ωq\_\{\\omega\}=\\hat\{k\}\_\{\\omega\}andqt\+h⋆=k\(Ut\+h⋆\)q^\{\\star\}\_\{t\+h\}=k\(U^\{\\star\}\_\{t\+h\}\), meaning for a fixed target, we aim to stabilize around a constantq⋆q^\{\\star\}\.
All tasks shareH=5H=5latent steps \(ten native steps\), score all five rollout states, and initialize with a nominal control actionaref=2a^\{\\mathrm\{ref\}\}=2\. The tracking\-only latent configuration uses\(α,λa\)=\(1,0\)\(\\alpha,\\lambda\_\{a\}\)=\(1,0\)\. The native\-reward comparison uses\(0\.05,0\.05\)\(0\.05,0\.05\)for latent\-L2L^\{2\}and\(β,λa\)=\(1,0\.05\)\(\\beta,\\lambda\_\{a\}\)=\(1,0\.05\)for KE where alpha is a hyperparameter, tuned to best approximate the scale of the true reference penalty\. For the three\-signal aperiodic suite we use\(β,λa\)=\(1,0\)\(\\beta,\\lambda\_\{a\}\)=\(1,0\)andTKE=0\.004T\_\{\\mathrm\{KE\}\}=0\.004\. Since no other cost term is active, this is equivalent to having MPPI temperature\(β,T\)=\(25,0\.1\)\(\\beta,T\)=\(25,0\.1\)\. The conventional latent\-L2L^\{2\}baseline retainsT=0\.1T=0\.1\.
As written, equation \([5](https://arxiv.org/html/2607.21644#S3.E5)\) implicitly assumes that Euclidean displacement in the learned representation ranks candidate actions like physical field error\. While this kind of joint\-embedding training naturally seems to encourage predictable, action\-conditioned features, it does not directly enforce isometry\. Equation \([6](https://arxiv.org/html/2607.21644#S3.E6)\) intentionally asks less of the representation\. The learned probe estimates the value of one observable making the comparison to the target downstream, and in physically calibrated units\.
## 4Experimental Protocol
### 4\.1Offline data and optimization
The offline training dataset contains 6930 trajectories, each consisting of 200 velocity fields and 199 corresponding actions\. We use an≈92/8\\approx 92/8train validation split, with 6306 trajectories assigned to training and 624 left for validation\. Candidate actions are explicitly constrained to the range\[0,4\]\[0,4\]\. The training dataset combines exploratory initial conditions, counterfactual branches where candidate signals are rolled out before splitting into several potential cases, and piece\-wise constant samples\. The evaluation targets are the control gym benchmark reference and three custom aperiodic pulse sequences\. The established aperiodic sample triplet is specifically audited for distance from the training dataset\. \(Section[5\.1\.2](https://arxiv.org/html/2607.21644#S5.SS1.SSS2)\)\.
Figure 3:Composition of the offline training corpus, where \(a\) denotes the categories of control actions, \(b\) the dataset composition, and \(c\) the initial conditions\.Our model is first trained for 50,000 optimizer steps with a batch size 128, learning rate3×10−43\\times 10^\{\-4\}, weight decay10−210^\{\-2\}, a 2000\-step warmup, cosine decay\. We then fine tune for 50,000 steps at a reduced peak learning rate of6×10−56\\times 10^\{\-5\}and a half\-cosine decay to zero\. The resulting model is the single world model behind every result in this paper\. Model training curves are shown in Figure[4](https://arxiv.org/html/2607.21644#S4.F4)\. For all closed\-loop control results, we present an average taken over ten spatially constant random initial velocity conditions sampled on\(−5,5\)\(\-5,5\)\. For the published\-benchmark reward comparison, we extend this protocol to 50 random initial conditions to match the published count in Control Gym\. All aggregate result figures depict 95% CI unless stated otherwise\.
Figure 4:Training and validation diagnostics over the full 0–100k schedule\. The action decoder remains accurate, rollout losses continue to decrease, and the validation latent effective rank grows throughout the continuation\.
## 5Goal\-Agnostic Control Results
We organize results by control objective at deployment\. We first consider a direct comparison of control discovery using raw latent\-L2L^\{2\}versus observable KE costs\. Between examples, only the cost function and its scalar weight\(s\) are changed\. We report reconstructed velocity\-field RMSE independently of the planner cost so that it can be demonstrated that KE tracking indeed establishes improved flow tracking in state space\.
### 5\.1Time\-dependent target trajectories
#### 5\.1\.1PDE Control Gym benchmark
We first sweep the KE scaleβ\\betawith the benchmark action\-reference penalty fixed at the papers chosenλa=0\.05\\lambda\_\{a\}=0\.05on ten matched clean episodes\. Figure[5](https://arxiv.org/html/2607.21644#S5.F5)shows a smooth control\-effort frontier\. The latent\-L2L^\{2\}convention lies inside this frontier: at the reward\-oriented choiceβ=1\\beta=1, KE planning improves both mean reward \(−11\.78\-11\.78versus−12\.78\-12\.78\) and late field RMSE \(0\.06990\.0699versus0\.07710\.0771\)\. Atβ=100\\beta=100, it reaches late field RMSE0\.01110\.0111,46%46\\%below even the latent tracking\-only configuration’s0\.02060\.0206, despite fighting the action penalty\.
Figure 5:Clean linear\-reference KE cost\-scale screen over ten matched episodes \(mean and 95% CI\)\. The action\-reference weight remains0\.050\.05\. Dashed lines and blue bands show the matched latent\-L2L^\{2\}mean and 95% CI\. The circledβ=1\\beta=1operating point is used for the 50\-episode benchmark comparison\.Table 1:Clean linear\-reference results\. Reward is higher when less negative; parentheses give the planner tracking scale and action\-reference weight\. Reward intervals are 95% CIs\. The first four rows use ten matched episodes; the reward\-oriented rows use the 50\-episode benchmark protocol\.Controllerepisodesrewardlate field RMSEtarget replay \(oracle\)10−18\.192±1\.367\-18\.192\\pm 1\.3670\.00390constant2\.52\.510−10\.877±1\.410\-10\.877\\pm 1\.4100\.05441latent\-L2L^\{2\}MPPI\(1,0\)\(1,0\)10−22\.011±1\.371\-22\.011\\pm 1\.3710\.02057KE\-probe MPPI\(100,0\.05\)\(100,0\.05\)10−21\.228±2\.005\-21\.228\\pm 2\.0050\.01109latent\-L2L^\{2\}MPPI\(0\.05,0\.05\)\(0\.05,0\.05\)50−12\.084±0\.857\-12\.084\\pm 0\.8570\.07650KE\-probe MPPI\(1,0\.05\)\(1,0\.05\)50−10\.901±0\.911\\mathbf\{\-10\.901\\pm 0\.911\}0\.06917The 50\-episode confirmation uses seeds 0–49 and reruns the latent baseline on the identical seeds\. KE improves reward by1\.18±0\.201\.18\\pm 0\.20paired points \(95% CI\), with higher reward in 49 of 50 episodes, and lowers late field RMSE in all 50\.
Table 2:Native\-reward comparison on the default linear target\. Published values are PDE Control Gym Table 3 means over 50 episodes\[[4](https://arxiv.org/html/2607.21644#bib.bib1)\]; ours is the reward\-matched mean over 50 seeded episodes\. Published agents act every solver step, whereas our macro action is held for two steps, so the ranking is contextual\.The observable\-aligned controller improves native reward by6\.9286\.928relative to the published SAC anchor and lies2\.9702\.970behind the published model\-based optimizer\. Published agents act every solver step whereas our action is held for two steps, so this ranking is not final\. Nevertheless, we show that a world model trained without any reward supervision or field reconstruction can achieve competitive performance on an arbitrary control task\.
#### 5\.1\.2Three distinct aperiodic targets
We next replace the benchmark reference by three deterministic aperiodic maneuvers composed of Gaussian pulses with unequal timing, sign, amplitude, and width\. The signals consist of an up down up triplet, a staggered multiscale sequence, and an alternating asymmetric sequence\. All remain within the actuator range\[0,4\]\[0,4\], and their pairwise action RMSEs range from0\.8440\.844to0\.9930\.993\. The simplest of the three signals can be shown as provably outside of the training dataset, by comparing it against all 6306 training action sequences, we can show a nearest\-training\-trajectory action RMSE of0\.3110\.311\(Figure[6](https://arxiv.org/html/2607.21644#S5.F6)\)\.
Figure 6:Action\-space out\-of\-distribution audit for the aperiodic triplet\. Left: the frozen triplet schedule and its nearest training action trace \(RMSE0\.3110\.311\)\. Right: empirical CDF of RMSE distances from the pulse and from a period\-0\.10\.1cosine comparator to all 6,306 training traces; dotted lines mark the nearest\-trace distances\.Here the controlled responses occupy a narrow family on which KE and field error move closely together\. As Figure[7](https://arxiv.org/html/2607.21644#S5.F7)and Table[3](https://arxiv.org/html/2607.21644#S5.T3)show, KE\-probe planning drastically improves upon, and corrects nearly all of the phase lag observed when using latent\-L2L^\{2\}planning\. Across the\(3×10\)\(3\\times 10\)30 matched episodes, it lowers late field RMSE by53\.0%53\.0\\%\(0\.022040\.02204versus0\.046870\.04687\), demonstrating a measurable improvement across all 30 episodes\. The individual late\-RMSE reductions are58\.2%58\.2\\%,49\.9%49\.9\\%, and50\.6%50\.6\\%for the triplet, staggered, and alternating signals, shown in Figure[7](https://arxiv.org/html/2607.21644#S5.F7)\.
Figure 7:Clean tracking of three aperiodic targets over ten matched initial conditions per signal \(mean with 95% confidence bands\)\. Columns show the triplet, staggered multiscale, and alternating asymmetric targets\. The top row compares target and achieved KE after the first ten native steps; the bottom row compares the target\-generating and executed boundary controls\. Target controls are shown only for evaluation and are not supplied to MPPI\. KE\-probe planning follows their turning points, whereas latent\-L2L^\{2\}planning exhibits phase lag\.Table 3:Three aperiodic targets over ten seeded initial conditions each with clean observations\. The pooled rows weight all 30 episodes equally\.The two MPPI costs have essentially equal pooled native reward \(−15\.245\-15\.245versus−15\.084\-15\.084\)\. KE planning uses10\.6%10\.6\\%greater action total variation, consistent with the sharper control reversals visible in the bottom row of Figure[7](https://arxiv.org/html/2607.21644#S5.F7)\. The gain is therefore specifically a physical\-tracking result, not a reward claim\.
Purely for the sake of visualizing the field agreement, we reconstructed the first aperiodic signal and apply both the generating \(oracle\) and learned \(JEPA\) control, the results are shown in Figure[8](https://arxiv.org/html/2607.21644#S5.F8)\.
Figure 8:Aperiodic pulse reconstruction for case A: Up\-down\-up triplet, where both the learned and true trajectory can be shown to bear significant similarity, reflecting the accuracy with which the generating control was able to be recovered\.
### 5\.2Fixed kinetic\-energy regulation
Kinetic energy per unit mass is defined as
k\(U\)=12HW∑i,j\(uij2\+vij2\)\.k\(U\)=\\frac\{1\}\{2HW\}\\sum\_\{i,j\}\\bigl\(u\_\{ij\}^\{2\}\+v\_\{ij\}^\{2\}\\bigr\)\.\(7\)A ridge map fitted on frozen predictor\-rollout latents, without changing any JEPA weight, achieves held\-outR2=0\.9891R^\{2\}=0\.9891\. The trajectory\-tracking experiments above use this same readout with a time\-varying target\. Here we hold the requested KE fixed to test scalar regulation\.
Figure 9:Derived\-quantity control with a rollout\-latent kinetic\-energy probe\. From ten random initial conditions, MPPI drives kinetic energy toward the scalar targetk⋆=0\.232k^\{\\star\}=0\.232\(left\) and converges to a mean action of2\.492\.49, close to the true equilibrium action2\.52\.5\(right\)\.Over the last 50 native steps, the controller has2\.66%2\.66\\%mean pointwise relative error, while the constant action\-reference at2\.02\.0has37\.7%37\.7\\%error\. Our observable\-aligned controller is shown to match the generating constant while discovering the holding action from the scalar energy target alone\. Its mean steady action is2\.491±0\.0442\.491\\pm 0\.044across episodes, closely recovering the target equilibrium action of2\.52\.5\.
### 5\.3Limitation: probe fragility under corrupted observations
The observable\-aligned tracking gains above are full\-state observation results\. This qualification is important: the kinetic energy probe is demonstrably unresilient to corrupted measurements as implemented\. We diagnose that failure using the fixed\-KE task\. Gaussian noise is drawn per step with standard deviation equal to55,1010, or15%15\\%of each velocity channel’s recent RMS; fixed random\-pixel masks remove55,1515, or25%25\\%of pixels\. Corruptions affect only the controller’s measurement, while the simulator and target remain clean\. Each cell uses the same ten seeds\.
The fixed\-field latent goal provides context rather than a headline robustness claim: its steady RMSE rises from0\.0220\.022clean to0\.0820\.082at15%15\\%noise\. The probe\-defined KE interface is much more brittle\. Its clean2\.7%2\.7\\%relative error jumps to4646–64%64\\%under noise and6060–85%85\\%under missing pixels, with a large steady\-action bias\. The ridge probe is fit on clean rollout latents, and corrupted measurements evidently shift the encoded state off the manifold on which the probe is calibrated\. Robustness is therefore task dependent, while the frozen dynamics can accept masked or incomplete measurements, the learned downstream readout must also be calibrated\. As such, we do not claim robustness to partial or corrupted observations for the KE controller\. A short corruption diagnostic is deferred to Appendix[B](https://arxiv.org/html/2607.21644#A2)\.
## 6Discussion
The experiments demonstrate the potential for JEPA world models to serve as goal\-agnostic, flexible controllers for PDE problems\. We show that the learned latent representation need not itself be the task metric, as JEPA can be relied upon primarily to supply action\-conditioned state evolution\. For our use case, where kinetic energy determines state exactly, a small readout can lead to a landscape where the control objective is better defined\. This separation of task ant latent representation preserves the claim of goal independence while avoiding the stronger and unnecessary assumption that raw latent Euclidean distance is metric preserving\.
Changing the cost from minimizingL2L^\{2\}distance to KE moves the same predictor from−12\.08\-12\.08to−10\.90\-10\.90mean reward in the 50\-episode comparison\. Across each of the aperiodic signals, the KE objective substantially improves control discovery, removing a phase lag present under latent\-L2L^\{2\}and halves late field error in all 30 matched episodes\. These gains indicate that the earlier bottleneck was the geometry of the cost function at least as much as the structure of the world model itself\.
Across a broad class of problems, kinetic energy is not to be considered a substitute for physical field distance\. We infer that our approacah succeeds here because a single scalar boundary actuator \(confined to positive values only\!\) restricts the evaluated trajectories to a narrow response family where the map from state to kinetic energy is actually injective\. Velocity\-field RMSE is therefore retained as an independent evaluation metric, and the result should not be extrapolated to multi\-actuator or topologically distinct flows without further observables\. We do however see this as potentially positive evidence toward the claim that for well understood systems, controlling a carefully selected observable which is known to guarantee unique continuation may prove a generalizable approach to overcoming the shortcomings of L2 distance\.
Our second proposed limitation is the requirement for full state observation\. Since our probe is fitted on clean rollout latents it fails when noise or missing pixels move the current\-state encoding even slightly away from the true embedding\. Our observable\-aligned planning in its current form thus trades an uncalibrated latent metric for an explicit readout that must itself be made robust\. Training readouts on corrupted latents, adding uncertainty quantification or other robustness measures offer avenues for natural extension\.
Every experiment still uses the same five\-step latent rollout\. No task\-wise horizon tuning is needed for the reported results, although longer horizons could alter individual operating points\. The central comparison is unchanged: the encoder, predictor, and MPPI optimizer are held fixed while only the deployment\-time cost interface changes\.
## 7Related Work
Planning in a learned latent space predates the introduction of joint\-embedding architectures\. For example, Embed\-to\-Control learns locally linear latent dynamics for control from raw images\[[34](https://arxiv.org/html/2607.21644#bib.bib12)\]; World Models\[[12](https://arxiv.org/html/2607.21644#bib.bib13)\]and PlaNet\[[14](https://arxiv.org/html/2607.21644#bib.bib14)\]either plan or learn behaviors inside learned latent dynamics, with Dreamer extending to latent imagination\[[13](https://arxiv.org/html/2607.21644#bib.bib15)\]and PETS establishing sampling\-based MPC over learned probabilistic dynamics\[[8](https://arxiv.org/html/2607.21644#bib.bib16)\]\. Closest to our work are TD\-MPC and TD\-MPC2, which run MPPI over a reconstruction\-free latent dynamics model trained with a latent consistency loss\[[16](https://arxiv.org/html/2607.21644#bib.bib17),[15](https://arxiv.org/html/2607.21644#bib.bib18)\], however in both cases latent spaces are shaped by task rewards and value learning, whereas a joint\-embedding approahc like ours is trained explicitly with no task signal at all\. MuZero also plans in a value\-equivalent latent space without reconstruction\[[30](https://arxiv.org/html/2607.21644#bib.bib19)\]\. Our contribution transfers this latent\-planning recipe to PDE dynamics while keeping the representation goal\-agnostic\.
JEPAs learn by predicting target embeddings rather than reconstructing inputs\[[22](https://arxiv.org/html/2607.21644#bib.bib20),[2](https://arxiv.org/html/2607.21644#bib.bib2)\], generally combining masked prediction with vision transformers\[[9](https://arxiv.org/html/2607.21644#bib.bib22),[17](https://arxiv.org/html/2607.21644#bib.bib24)\]and EMA target encoders\[[11](https://arxiv.org/html/2607.21644#bib.bib23)\]\. For end to end models such as ours, additional anti\-collapse measures must be taken, VICReg supplies explicit variance and covariance regularization against collapse\[[3](https://arxiv.org/html/2607.21644#bib.bib3)\]\. DINO\-WM demonstrates that pretrained non\-generative visual features can support zero\-shot planning\[[39](https://arxiv.org/html/2607.21644#bib.bib5)\]; V\-JEPA 2 post\-trains an action\-conditioned predictor on robot interaction data and plans manipulation from image goals\[[1](https://arxiv.org/html/2607.21644#bib.bib21)\]; LeWorldModel learns an end\-to\-end JEPA for manipulation, navigation, and locomotion tasks\[[24](https://arxiv.org/html/2607.21644#bib.bib8)\]; SkyJEPA extends JEPA world models to long\-horizon quadrotor control\[[29](https://arxiv.org/html/2607.21644#bib.bib11)\]\.Delta\-JEPA target the complementary control problem of action\-insensitive latent dynamics by decoding actions from latent differences\[[38](https://arxiv.org/html/2607.21644#bib.bib7)\]\. While existing work establish JEPA world modeling as a promising avenue for visual and robotic domains, to our knowledge, ours is the first to apply a JEPA world model to closed\-loop PDE control explicitly\.
Learned PDE surrogates commonly evolve either full fields or compressed states\. Differentiable\-solver control\[[18](https://arxiv.org/html/2607.21644#bib.bib25)\], operator\-learning control\[[19](https://arxiv.org/html/2607.21644#bib.bib26)\], PINN\-based optimal control\[[26](https://arxiv.org/html/2607.21644#bib.bib27)\], and diffusion\-based control\[[35](https://arxiv.org/html/2607.21644#bib.bib28)\]optimize actions through a surrogate coupled to a full state\-space objective, while neural operators have shown promise in learning the gain computations of PDE backstepping directly\[[5](https://arxiv.org/html/2607.21644#bib.bib29)\]\. LE\-PDE proposes a latent\-evolution approach that targets both accelerated simulation and inverse optimization\[[37](https://arxiv.org/html/2607.21644#bib.bib6)\]\. Closest in spirit to our setting are Koopman and reduced\-order methods performing control in a learned low\-dimensional spaces\. Examples of this include deep Koopman embeddings\[[23](https://arxiv.org/html/2607.21644#bib.bib31)\], Koopman\-based MPC\[[21](https://arxiv.org/html/2607.21644#bib.bib32),[27](https://arxiv.org/html/2607.21644#bib.bib33)\], SINDy\-MPC\[[20](https://arxiv.org/html/2607.21644#bib.bib34)\], latent MPC of unsteady fluid flows\[[25](https://arxiv.org/html/2607.21644#bib.bib30)\], and deep\-ROM latent feedback control\[[31](https://arxiv.org/html/2607.21644#bib.bib35)\]\. Our focus differs from existing latent evolution works by adopting the JEPA / World Modeling framework, removing the field decoder entirely and fully separating the learned dynamics from the deployment goal\.
## 8Conclusion
Goal\-agnostic joint\-embedding control separates learning action\-conditioned dynamics from control objectives\. From our observations, the strongest results are achieved when the control objective can be described deterministically in terms of an explicit physical observable rather than raw latent distance\. We find that a rank\-one KE probe improves performance for the linear\-reference reward, field RMSE, reduces time\-dependent tracking error by53%53\\%across three aperiodic targets relative to latent\-L2L^\{2\}alone\. The same checkpoint and short\-horizon planner retain fixed\-field stabilization as a separate goal interface\.
While KE results are not able to be generalized across settings, and the clean\-calibrated probe is brittle under corrupted observations, we believe a properly prescribed observable, if capable of being faithfully reconstructed from the latent space could transfer the benefits found in our work\. Our experiments show that reconstruction\-free latent dynamics are nonetheless a powerful substrate for control, while a carefully selected physics probe can supply a better control objective, even if state control is the primary goal\. The practical lesson is not that every PDE should be controlled by energy, but that JEPA world models are capable of achieving competitive control results for PDE in a goal agnostic setting, and furthermore that if available, control through a calibrated observable \(provided it provides uniqueness\) can improve state space tracking\.
## Appendix AHyperparameter Summary
Table 4:Principal model and planner settings\.
## Appendix BCorrupted\-Observation Diagnostic
Figure 10:Secondary corruption diagnostic over ten matched episodes\. Fixed\-field stabilization remains usable across most conditions, whereas the clean\-calibrated kinetic energy probe fails under both measurement noise and missing pixels \(mean±\\pm95% CI\)\.
## References
- \[1\]M\. Assran, A\. Bardes, D\. Fan, Q\. Garrido, R\. Howes,et al\.\(2025\)V\-JEPA 2: self\-supervised video models enable understanding, prediction and planning\.arXiv preprint arXiv:2506\.09985\.External Links:[Link](https://arxiv.org/abs/2506.09985)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p4.1),[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[2\]M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, Y\. LeCun, and N\. Ballas\(2023\)Self\-supervised learning from images with a joint\-embedding predictive architecture\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,External Links:[Link](https://arxiv.org/abs/2301.08243)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p2.1),[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[3\]A\. Bardes, J\. Ponce, and Y\. LeCun\(2022\)VICReg: variance\-invariance\-covariance regularization for self\-supervised learning\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2105.04906)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p5.1),[§3\.1](https://arxiv.org/html/2607.21644#S3.SS1.p3.1),[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[4\]L\. Bhan, Y\. Bian, M\. Krstic, and Y\. Shi\(2024\)PDE Control Gym: a benchmark for data\-driven boundary control of partial differential equations\.arXiv preprint arXiv:2405\.11401\.External Links:[Link](https://arxiv.org/abs/2405.11401)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p5.1),[§2\.3](https://arxiv.org/html/2607.21644#S2.SS3.p1.8),[Table 2](https://arxiv.org/html/2607.21644#S5.T2),[Table 2](https://arxiv.org/html/2607.21644#S5.T2.9.2)\.
- \[5\]L\. Bhan, Y\. Shi, and M\. Krstic\(2024\)Neural operators for bypassing gain and control computations in PDE backstepping\.IEEE Transactions on Automatic Control69,pp\. 5310–5325\.External Links:[Link](https://arxiv.org/abs/2302.14265)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[6\]S\. L\. Brunton and B\. R\. Noack\(2015\)Closed\-loop turbulence control: progress and challenges\.Applied Mechanics Reviews67\(5\),pp\. 050801\.External Links:[Document](https://dx.doi.org/10.1115/1.4031175)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1)\.
- \[7\]M\. A\. Bucci, O\. Semeraro, A\. Allauzen, G\. Wisniewski, L\. Cordier, and L\. Mathelin\(2019\)Control of chaotic systems by deep reinforcement learning\.Proceedings of the Royal Society A475\(2231\),pp\. 20190351\.External Links:[Document](https://dx.doi.org/10.1098/rspa.2019.0351)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1)\.
- \[8\]K\. Chua, R\. Calandra, R\. McAllister, and S\. Levine\(2018\)Deep reinforcement learning in a handful of trials using probabilistic dynamics models\.InAdvances in Neural Information Processing Systems,External Links:[Link](https://arxiv.org/abs/1805.12114)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[9\]A\. Dosovitskiy, L\. Beyer, A\. Kolesnikov, D\. Weissenborn, X\. Zhai, T\. Unterthiner, M\. Dehghani, M\. Minderer, G\. Heigold, S\. Gelly, J\. Uszkoreit, and N\. Houlsby\(2021\)An image is worth 16x16 words: transformers for image recognition at scale\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2010.11929)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[10\]A\. Farahmand, S\. Nabi, and D\. N\. Nikovski\(2017\)Deep reinforcement learning for partial differential equation control\.In2017 American Control Conference,pp\. 3120–3127\.External Links:[Document](https://dx.doi.org/10.23919/ACC.2017.7963427)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1)\.
- \[11\]J\. Grill, F\. Strub, F\. Altché, C\. Tallec, P\. H\. Richemond, E\. Buchatskaya, C\. Doersch, B\. A\. Pires, Z\. D\. Guo, M\. G\. Azar, B\. Piot, K\. Kavukcuoglu, R\. Munos, and M\. Valko\(2020\)Bootstrap your own latent: a new approach to self\-supervised learning\.InAdvances in Neural Information Processing Systems,External Links:[Link](https://arxiv.org/abs/2006.07733)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[12\]D\. Ha and J\. Schmidhuber\(2018\)World models\.arXiv preprint arXiv:1803\.10122\.External Links:[Link](https://arxiv.org/abs/1803.10122)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[13\]D\. Hafner, T\. Lillicrap, J\. Ba, and M\. Norouzi\(2020\)Dream to control: learning behaviors by latent imagination\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/1912.01603)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[14\]D\. Hafner, T\. Lillicrap, I\. Fischer, R\. Villegas, D\. Ha, H\. Lee, and J\. Davidson\(2019\)Learning latent dynamics for planning from pixels\.InInternational Conference on Machine Learning,External Links:[Link](https://arxiv.org/abs/1811.04551)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[15\]N\. Hansen, H\. Su, and X\. Wang\(2024\)TD\-MPC2: scalable, robust world models for continuous control\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2310.16828)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[16\]N\. Hansen, X\. Wang, and H\. Su\(2022\)Temporal difference learning for model predictive control\.InInternational Conference on Machine Learning,External Links:[Link](https://arxiv.org/abs/2203.04955)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[17\]K\. He, X\. Chen, S\. Xie, Y\. Li, P\. Dollár, and R\. Girshick\(2022\)Masked autoencoders are scalable vision learners\.InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,External Links:[Link](https://arxiv.org/abs/2111.06377)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[18\]P\. Holl, V\. Koltun, and N\. Thuerey\(2020\)Learning to control PDEs with differentiable physics\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2001.07457)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1),[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[19\]R\. Hwang, J\. Y\. Lee, J\. Y\. Shin, and H\. J\. Hwang\(2022\)Solving PDE\-constrained control problems using operator learning\.InProceedings of the AAAI Conference on Artificial Intelligence,pp\. 4504–4512\.External Links:[Link](https://arxiv.org/abs/2111.04941)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1),[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[20\]E\. Kaiser, J\. N\. Kutz, and S\. L\. Brunton\(2018\)Sparse identification of nonlinear dynamics for model predictive control in the low\-data limit\.Proceedings of the Royal Society A474\(2219\),pp\. 20180335\.External Links:[Document](https://dx.doi.org/10.1098/rspa.2018.0335)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[21\]M\. Korda and I\. Mezić\(2018\)Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control\.Automatica93,pp\. 149–160\.External Links:[Document](https://dx.doi.org/10.1016/j.automatica.2018.03.046)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[22\]Y\. LeCun\(2022\)A path towards autonomous machine intelligence\.Note:OpenReview preprint, version 0\.9\.2External Links:[Link](https://openreview.net/forum?id=BZ5a1r-kVsf)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[23\]B\. Lusch, J\. N\. Kutz, and S\. L\. Brunton\(2018\)Deep learning for universal linear embeddings of nonlinear dynamics\.Nature Communications9,pp\. 4950\.External Links:[Document](https://dx.doi.org/10.1038/s41467-018-07210-0)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[24\]L\. Maes, Q\. Le Lidec, D\. Scieur, Y\. LeCun, and R\. Balestriero\(2026\)LeWorldModel: stable end\-to\-end joint\-embedding predictive architecture from pixels\.arXiv preprint arXiv:2603\.19312\.External Links:[Link](https://arxiv.org/abs/2603.19312)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p4.1),[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[25\]J\. Morton, A\. Jameson, M\. J\. Kochenderfer, and F\. Witherden\(2018\)Deep dynamical modeling and control of unsteady fluid flows\.InAdvances in Neural Information Processing Systems,External Links:[Link](https://arxiv.org/abs/1805.07472)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[26\]S\. Mowlavi and S\. Nabi\(2023\)Optimal control of PDEs using physics\-informed neural networks\.Journal of Computational Physics473,pp\. 111731\.External Links:[Link](https://arxiv.org/abs/2111.09880)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1),[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[27\]S\. Peitz and S\. Klus\(2019\)Koopman operator\-based model reduction for switched\-system control of PDEs\.Automatica106,pp\. 184–191\.External Links:[Document](https://dx.doi.org/10.1016/j.automatica.2019.05.016)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[28\]J\. Rabault, M\. Kuchta, A\. Jensen, U\. Réglade, and N\. Cerardi\(2019\)Artificial neural networks trained through deep reinforcement learning discover control strategies for active flow control\.Journal of Fluid Mechanics865,pp\. 281–302\.External Links:[Document](https://dx.doi.org/10.1017/jfm.2019.62)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1)\.
- \[29\]P\. Rao, W\. Zhang, R\. Balestriero, Y\. LeCun, and G\. Loianno\(2026\)SkyJEPA: learning long\-horizon world models for zero\-shot sim\-to\-real control of quadrotors\.External Links:2606\.23444,[Link](https://arxiv.org/abs/2606.23444)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p4.1),[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[30\]J\. Schrittwieser, I\. Antonoglou, T\. Hubert, K\. Simonyan, L\. Sifre, S\. Schmitt, A\. Guez, E\. Lockhart, D\. Hassabis, T\. Graepel, T\. Lillicrap, and D\. Silver\(2020\)Mastering Atari, Go, chess and shogi by planning with a learned model\.Nature588,pp\. 604–609\.External Links:[Document](https://dx.doi.org/10.1038/s41586-020-03051-4)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[31\]M\. Tomasetto, F\. Braghin, and A\. Manzoni\(2024\)Latent feedback control of distributed systems in multiple scenarios through deep learning\-based reduced order models\.arXiv preprint arXiv:2412\.09942\.External Links:[Link](https://arxiv.org/abs/2412.09942)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[32\]S\. Verma, G\. Novati, and P\. Koumoutsakos\(2018\)Efficient collective swimming by harnessing vortices through deep reinforcement learning\.Proceedings of the National Academy of Sciences115\(23\),pp\. 5849–5854\.External Links:[Document](https://dx.doi.org/10.1073/pnas.1800923115)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1)\.
- \[33\]Y\. Wang, O\. Bounou, G\. Zhou, R\. Balestriero, T\. G\. J\. Rudner, Y\. LeCun, and M\. Ren\(2026\)Temporal straightening for latent planning\.External Links:2603\.12231,[Link](https://arxiv.org/abs/2603.12231)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p3.1)\.
- \[34\]M\. Watter, J\. T\. Springenberg, J\. Boedecker, and M\. Riedmiller\(2015\)Embed to control: a locally linear latent dynamics model for control from raw images\.InAdvances in Neural Information Processing Systems,External Links:[Link](https://arxiv.org/abs/1506.07365)Cited by:[§7](https://arxiv.org/html/2607.21644#S7.p1.1)\.
- \[35\]L\. Wei, P\. Hu, R\. Feng, H\. Feng, Y\. Du, T\. Zhang, R\. Wang, Y\. Wang, Z\. Ma, and T\. Wu\(2024\)DiffPhyCon: a generative approach to control complex physical systems\.InAdvances in Neural Information Processing Systems,External Links:[Link](https://arxiv.org/abs/2407.06494)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1),[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[36\]G\. Williams, N\. Wagener, B\. Goldfain, P\. Drews, J\. M\. Rehg, B\. Boots, and E\. A\. Theodorou\(2017\)Information theoretic MPC for model\-based reinforcement learning\.In2017 IEEE International Conference on Robotics and Automation,pp\. 1714–1721\.External Links:[Document](https://dx.doi.org/10.1109/ICRA.2017.7989202),[Link](https://ieeexplore.ieee.org/document/7989202)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p5.1),[§3\.2](https://arxiv.org/html/2607.21644#S3.SS2.p1.2)\.
- \[37\]T\. Wu, T\. Maruyama, and J\. Leskovec\(2022\)Learning to accelerate partial differential equations via latent global evolution\.arXiv preprint arXiv:2206\.07681\.External Links:[Link](https://arxiv.org/abs/2206.07681)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p1.1),[§1](https://arxiv.org/html/2607.21644#S1.p2.1),[§7](https://arxiv.org/html/2607.21644#S7.p3.1)\.
- \[38\]Z\. Zhang, Y\. Wang, Z\. Guan, Y\. Yang, B\. Shi, T\. Zong, H\. Yi, G\. Chao, X\. Chen, T\. Yang, C\. Bao, T\. Yu, J\. Zhou, and J\. Xu\(2026\)Delta\-JEPA: learning action\-sensitive world models via latent difference decoding\.arXiv preprint arXiv:2606\.31232\.External Links:[Link](https://arxiv.org/abs/2606.31232)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p4.1),[§1](https://arxiv.org/html/2607.21644#S1.p5.1),[§3\.1](https://arxiv.org/html/2607.21644#S3.SS1.p3.2),[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.
- \[39\]G\. Zhou, H\. Pan, Y\. LeCun, and L\. Pinto\(2024\)DINO\-WM: world models on pre\-trained visual features enable zero\-shot planning\.arXiv preprint arXiv:2411\.04983\.External Links:[Link](https://arxiv.org/abs/2411.04983)Cited by:[§1](https://arxiv.org/html/2607.21644#S1.p2.1),[§1](https://arxiv.org/html/2607.21644#S1.p4.1),[§7](https://arxiv.org/html/2607.21644#S7.p2.1)\.Similar Articles
Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
Proposes temporal-distance JEPA (TD-JEPA) which mines directed temporal cost from offline trajectories to improve latent world model predictive control, achieving higher success rates on robotic environments.
Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding
Delta-JEPA introduces a reconstruction-free world model that augments latent forward prediction with a Latent Difference Action Decoder to prevent collapse and improve action-sensitivity, achieving better planning performance on visual continuous-control tasks.
AeroJEPA: Learning Semantic Latent Representations for Scalable 3D Aerodynamic Field Modeling
This paper introduces AeroJEPA, a Joint-Embedding Predictive Architecture for scalable 3D aerodynamic field modeling. It addresses limitations in current surrogate models by predicting semantic latent representations of flow fields, enabling efficient high-fidelity analysis and design optimization.
JEPA for AI-Native 6G: Predictive Representations and Open Challenges
This paper presents a tutorial on using Joint-Embedding Predictive Architecture (JEPA) for self-supervised learning in 6G networks, along with a beam management case study and open challenges.
Prediction and control with temporal segment models
OpenAI introduces a method for learning complex nonlinear system dynamics using deep generative models over temporal segments, enabling stable long-horizon predictions and differentiable trajectory optimization for model-based control.