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

Summary

The goal is an independent JEPA framework that trains a frozen 2D ViT latent model offline, without using reward signals or downstream targets. The model is then reused for PDE control with an MPPI controller. Replacing the L² target in latent space with a kinetic energy probe function significantly improves tracking performance in Navier–Stokes comparisons.

Why it matters

Decoupling world model learning from specific control objectives allows the same frozen model to be reused for different tasks without retraining. This can reduce iteration costs in fluid dynamics and other PDE-based control scenarios. The reported R²=0.989 value for held-out trajectories and the 53% RMSE reduction achieved with the KE probe function demonstrate that this approach is not merely a theoretical curiosity but a practical planning tool.

Highlights

• Combining a frozen latent model with MPPI reproduces held-out trajectories with R²=0.989 accuracy, requiring zero retraining
• Planning with the KE probe function reduces the final field RMSE by 53% (0.0220 vs 0.0469) on three held-out and non-periodic targets
• The latent probe function is explicitly vulnerable to measurement noise and missing pixels; this is a known and unresolved weakness

Critical take

These gains depend on the existence of a calibrated observable quantity that guarantees a unique continuation. This assumption is a non-trivial precondition that limits the framework's general validity far beyond the 'target-agnostic' scope. Additionally, the accepted vulnerability to sensor noise and missing pixels poses a practical barrier for real-world deployment scenarios. However, the article does not offer a reduction method beyond the single KE probe function example.