# Toward Goal-Agnostic Joint-Embedding Predictive Control of Partial Differential Equations

arXiv: https://arxiv.org/abs/2607.21644 · HTML: https://mlsanity.com/paper/2607.21644/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 40 references that are in the map.
- First posted on arXiv 21 Jul 2026.

## Builds on (most cited first, 30 of 32)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Deep learning for universal linear embeddings of nonlinear dynamics](https://mlsanity.com/paper/1712.09707/) (2017)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning](https://mlsanity.com/paper/2411.04983/) (2024)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels](https://mlsanity.com/paper/2603.19312/) (2026)
- [Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control](https://mlsanity.com/paper/1611.03537/) (2016)
- [Sparse identification of nonlinear dynamics for model predictive control in the low-data limit](https://mlsanity.com/paper/1711.05501/) (2017)
- [Learning to Control PDEs with Differentiable Physics](https://mlsanity.com/paper/2001.07457/) (2020)
- [Efficient collective swimming by harnessing vortices through deep reinforcement learning](https://mlsanity.com/paper/1802.02674/) (2018)
- [Deep Dynamical Modeling and Control of Unsteady Fluid Flows](https://mlsanity.com/paper/1805.07472/) (2018)
- [Optimal control of PDEs using physics-informed neural networks](https://mlsanity.com/paper/2111.09880/) (2021)
- [Solving PDE-constrained Control Problems Using Operator Learning](https://mlsanity.com/paper/2111.04941/) (2021)
- [Learning to Accelerate Partial Differential Equations via Latent Global Evolution](https://mlsanity.com/paper/2206.07681/) (2022)
- [Koopman operator-based model reduction for switched-system control of PDEs](https://mlsanity.com/paper/1710.06759/) (2017)
- [Control of chaotic systems by Deep Reinforcement Learning](https://mlsanity.com/paper/1906.07672/) (2019)
- [Neural Operators for Bypassing Gain and Control Computations in PDE Backstepping](https://mlsanity.com/paper/2302.14265/) (2023)
- [DiffPhyCon: A Generative Approach to Control Complex Physical Systems](https://mlsanity.com/paper/2407.06494/) (2024)
- [Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding](https://mlsanity.com/paper/2606.31232/) (2026)

## Abstract (arXiv)

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 vision-transformer (ViT) and action-conditioned latent dynamics are trained offline without a reward or downstream goal, before being frozen and reused by a model-predictive path integral (MPPI) controller. We minimize a control objective in the latent space, initially expressed via the L2 distance and additionally illustrate the benefit of recasting the control objective in terms of an explicit physical observable when available. By instead minimizing the tracking error for a learned linear kinetic-energy (KE) probe on the frozen latent-state rollouts, we demonstrate the ability to reproduce the control of held-out trajectories with R2=0.989, while requiring no change to the underlying world model. For a controlled 2D Navier--Stokes benchmark, using a KE-probe within MPPI planning improves the mean native reward from -12.08±0.86 for latent-L2 tracking to -10.90±0.91 (95% CI), all while lowering last-quarter velocity-field RMSE from 0.0765 to 0.0692. Across three intentionally withheld, dissimilar, aperiodic targets, KE planning lowers late field RMSE by 53% relative to latent-L2 planning (0.0220 versus 0.0469), winning across 30 paired comparisons. The same frozen model also supports stabilization around a steady-state configuration via direct regulation of KE, achieving 2.7% mean relative error. While the latent probe proves brittle to measurement noise and missing pixels, our findings support the claim that latent dynamics can remain flexible and goal-agnostic, particularly when calibrated observables (granted they guarantee unique continuation) are a suitable objective for state control.