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

2026-07-21

AI Summary

This work addresses the challenge of achieving general and efficient control of partial differential equation (PDE) systems without task-specific objectives or reward signals. To this end, it proposes a goal-agnostic PDE control framework that combines an offline-trained, frozen Vision Transformer (ViT) encoder with an action-conditional latent dynamics model based on the Joint Embedding Predictive Architecture (JEPA), integrated with Model Predictive Path Integral (MPPI) control for online planning. Innovatively, the approach couples goal-agnostic latent space modeling with probes on physically meaningful observables—such as kinetic energy—enabling a single frozen world model to support diverse control tasks. Evaluated on Navier–Stokes benchmarks, the method substantially improves performance: kinetic-energy-probe-based planning raises the 50-episode average reward from −12.08 to −10.90 and reduces late-stage velocity field RMSE by 9.5%; across three unseen non-periodic targets, it cuts late-field RMSE by 53% and wins all 30 head-to-head trials; steady-state control achieves a 2.7% average relative error.

Abstract

We present a goal-agnostic control framework for partial differential equations (PDEs) built around a joint-embedding predictive architecture (JEPA). The small 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 ($L^2$) in the learned latent space. For a learned linear kinetic-energy (KE) probe on frozen latent rollouts we can reproduce held-out trajectories with $R^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\pm0.86$ for latent-$L^2$ planning to $-10.90\pm0.91$ (95\% CI), 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-$L^2$ planning ($0.0220$ versus $0.0469$), winning all 30 paired episodes. The same frozen model also supports controls targeting 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, 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

Problem

• goal-agnostic control
• partial differential equations
• latent dynamics
• predictive control
• state control

Innovation

• goal-agnostic control
• joint-embedding predictive architecture
• latent dynamics
• physical observables
• PDE control