JEPA framework controls PDEs with 53% less error using kinetic energy probe

Gallagher and Guglielmi present a goal-agnostic control framework for partial differential equations (PDEs) centered on a joint-embedding predictive architecture (JEPA). The system uses a small 2D Vision Transformer encoder paired with action-conditioned latent dynamics, all trained offline without any reward or downstream objective. The frozen latent model is then leveraged by a model-predictive path integral (MPPI) controller. The key insight is that when a control objective is available, it is far more effective to apply it to an explicit physical observable (provided injectivity holds) rather than minimizing raw Euclidean distance (L^2) in the learned latent space. A learned linear kinetic-energy (KE) probe on the frozen latent rollouts achieves R^2=0.989 for reproducing held-out trajectories, requiring no changes to the underlying world model.

On the PDE Control Gym's 2D Navier-Stokes benchmark, KE-probe planning significantly outperforms latent L^2 planning: it improves the matched 50-episode native reward from -12.08±0.86 to -10.90±0.91 (95% CI), and lowers the last-quarter velocity-field RMSE from 0.0765 to 0.0692. More strikingly, across three intentionally withheld, dissimilar, aperiodic targets, KE planning reduces late field RMSE by 53% relative to latent L^2 planning (0.0220 versus 0.0469), winning all 30 paired episodes. The same frozen model also successfully controls stabilization around a steady configuration via direct regulation of KE, achieving 2.7% mean relative error. While the latent probe is brittle to measurement noise and missing pixels, the results strongly support the claim that latent dynamics can remain both dynamic and goal-agnostic, with calibrated observables serving as a superior objective for state control.