🤖 machine learning

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

This paper introduces a goal-agnostic control framework for partial differential equations using a frozen joint-embedding predictive architecture, demonstrating that applying control objectives to explicit physical observables (like kinetic energy) rather than raw latent distances significantly improves trajectory matching, target tracking, and stabilization performance.

Original authors: Jonathan Gallagher, Roberto Guglielmi

Published 2026-07-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jonathan Gallagher, Roberto Guglielmi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to steer a massive, chaotic cloud of smoke or a swirling river of water. In the world of physics, these swirling patterns are described by equations called Partial Differential Equations (PDEs). Think of these equations as the ultimate rulebook for how fluids move, but the rulebook is so thick and complicated that trying to calculate the perfect move in real-time is like trying to solve a million-piece puzzle while the pieces are constantly changing shape. Scientists have been trying to build "autopilots" for these fluids using artificial intelligence, but most current methods are like students who memorize the answer key for one specific test. If the teacher changes the question slightly, the student fails because they never learned the underlying logic, only the specific solution.

To fix this, researchers are exploring a new way of learning called "Joint-Embedding Predictive Architectures" (JEPA). Instead of trying to memorize the exact picture of the fluid (which is hard and slow), a JEPA learns a "mental map" of how the fluid behaves. It's like learning the rules of a video game by watching how characters move, rather than trying to draw every single pixel of the game world. This allows the AI to predict what will happen next if you push a button, without needing to see the whole picture perfectly. The big question is: can this "mental map" be used to steer a real, messy fluid to a specific goal, even if the goal changes at the last second?

This paper introduces a new "goal-agnostic" control framework that uses a JEPA to steer a 2D fluid (simulating the famous Navier-Stokes equations) without ever needing to retrain the AI when the goal changes. The researchers trained a lightweight AI to understand how actions (like pushing on the edge of the fluid) change the fluid's future state, all while the AI was blind to any specific target. It learned the "physics" of the movement, not the destination. Once trained, this AI model was frozen and handed over to a planner (called MPPI) that could be told to aim for any goal, from following a wiggly line to keeping the fluid calm.

However, the team discovered a tricky surprise: just because the AI has a "mental map" doesn't mean the map is a perfect ruler. If you try to steer the fluid by measuring the distance between two points on this mental map (using a standard mathematical distance called L2), the AI gets confused. It's like trying to navigate a city using a map where the distance between "Home" and "School" is measured in "feels-like" units rather than actual miles; the AI might think it's close when it's actually far away. The paper argues against using this raw "mental distance" as the primary goal. Instead, the authors found that the AI performs much better when it is told to aim for a specific, real-world physical property, like the total "kinetic energy" (the amount of movement) of the fluid.

By attaching a simple "probe" (a tiny, fast calculator) to the AI's mental map, the system could translate the fluid's hidden state into a real-world number: its kinetic energy. The planner then used this number to steer the fluid. The results were impressive. When the team tested this on a benchmark where the fluid had to follow a specific path, the new method improved the score from -12.08 to -10.90 (where a higher, less negative number is better) and reduced the error in the fluid's speed by nearly 10%. Even more striking, when they tested the system on three completely new, unpredictable wiggly paths that the AI had never seen before, the "kinetic energy" method reduced the tracking error by 53% compared to the old method. It successfully followed the new targets without needing a single second of retraining.

The paper also showed that this same frozen AI could be used to stabilize the fluid, keeping it calm around a steady state by simply regulating its kinetic energy, achieving a mean relative error of just 2.7%. However, the authors are careful to note that this "kinetic energy probe" is a bit fragile. If the sensors measuring the fluid get noisy or if some data is missing (like pixels disappearing from a camera), the probe breaks down and the control fails. The paper suggests that while the AI's "mental map" is robust enough to handle missing data, the tool used to read the map needs to be calibrated for clean data.

In short, this research suggests that for controlling complex fluids, it's better to teach an AI the rules of the game and then let it aim for a clear, physical scoreboard (like total energy) rather than trying to make the AI's internal "feeling" of distance match reality perfectly. While the method isn't a magic bullet for every possible fluid problem (especially if the sensors are broken), it proves that a single, goal-independent AI model can be a flexible and powerful tool for steering complex physical systems, provided we give it the right kind of target to aim for.

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

Problem Statement

Controlling systems governed by Partial Differential Equations (PDEs), such as fluid dynamics, is challenging due to high-dimensional state spaces, nonlinear relationships, and the computational cost of numerical simulations required for evaluating control sequences. Existing methods face specific limitations:

Model-based control often requires problem-specific analysis or reduced-order models (ROMs) that may lack generality.
Learned surrogate models typically train on state-space forward simulation optimized for specific objectives, necessitating retraining when goals change and inheriting the difficulty of full state reconstruction.
Reinforcement Learning (RL) performs well on defined tasks but struggles with generalization due to the sensitivity of reward function selection.

The authors propose a framework that decouples the learning of system dynamics from the control objective, aiming for a "goal-agnostic" controller that can adapt to new tasks without retraining the underlying world model.

Methodology

The paper introduces a control framework built around a Joint-Embedding Predictive Architecture (JEPA) applied to PDEs. The core philosophy is to learn a latent world model that predicts how actions map between latent states, without reconstructing the physical state or using reward signals during training.

1. Goal-Agnostic World Model Training:

Architecture: A lightweight 2D Vision Transformer (ViT) encoder ( $E_\theta$ ) and an action-conditioned latent predictor ( $g_\phi$ ) are trained offline on state transitions.
Training Objective: The model is trained using masked observations (randomly masking 30–90% of the field) to predict future latent states. The loss function includes:
- Rollout Loss: Minimizing the distance between predicted and actual future latents.
- VICReg: Variance and covariance penalties to prevent representation collapse.
- Temporal Straightening: Encouraging latent trajectories to be direct rather than circuitous.
- Action Sensitivity: An auxiliary decoder predicts actions from latent displacements (inspired by Delta-JEPA) to ensure the latent space is sensitive to control inputs.
Key Constraint: No decoder maps latent states back to the physical field $U$ , and no reward function is used during training. The encoder and predictor are frozen after training.

2. Deployment and Control Strategy:

Model-Predictive Path Integral (MPPI): A frozen predictor is used within an MPPI controller to evaluate candidate action sequences.
Goal Interfaces: The framework supports three deployment modes without changing the world model weights:
1. Time-dependent tracking: Following a sequence of encoded states or physical observations.
2. Fixed-target stabilization: Maintaining a specific state configuration.
3. Kinetic Energy (KE) regulation: Stabilizing around a specific energy level.
The Observable-Aligned Controller: A critical innovation is the use of a lightweight linear probe ( $q_\omega$ $q_{ω}$ ) trained on frozen latent rollouts to estimate a physical observable (specifically, Kinetic Energy).
- Instead of minimizing raw Euclidean distance ( $L_2$ ) in the latent space (which may not correspond to physical error), the MPPI cost function minimizes the error between the predicted physical observable and the target observable.
- This approach relies on the assumption that the chosen observable (KE) guarantees unique continuation for the specific system constraints.

Key Contributions

First JEPA for Closed-Loop PDE Control: The authors present the first implementation of closed-loop PDE control using a joint-embedding latent world model that learns action-conditioned dynamics independently of the downstream task.
Observable-Aligned Planning: They introduce a method where a cheap linear probe estimates a physical scalar (Kinetic Energy) from latent predictions, and the controller optimizes tracking error in physical units rather than latent distance.
Goal Agnosticism via Frozen Models: A single frozen model checkpoint and a short (5-step) latent horizon successfully handle time-dependent tracking, fixed-field stabilization, and kinetic energy regulation without task-specific fine-tuning or retraining.
Demonstrated Performance Gains: On the PDE Control Gym 2D Navier–Stokes benchmark, the KE-probe approach outperforms latent- $L_2$ planning, improving native rewards and reducing field RMSE.

Experimental Results

The framework was evaluated on the PDE Control Gym 2D Navier–Stokes environment (unit square, $21 \times 21$ grid).

Benchmark Performance (50 episodes):
- Switching from latent- $L_2$ cost to KE-probe cost improved the mean native reward from −12.08 to −10.90.
- The last-quarter velocity-field RMSE decreased from 0.0765 to 0.0692.
- The KE-probe controller won 49 out of 50 paired episodes against the latent- $L_2$ baseline.
Aperiodic Targets:
- Tested against three distinct, withheld, aperiodic target signals.
- KE planning reduced late-field RMSE by 53% (0.0220 vs. 0.0469) compared to latent- $L_2$ planning, winning all 30 paired episodes.
- It successfully corrected phase lags observed in the latent- $L_2$ approach.
Stabilization:
- For fixed-field stabilization, the KE probe achieved a mean relative error of 2.7% in controlling the system to a steady configuration, recovering the equilibrium action with high precision.
Limitations:
- The KE probe is brittle to measurement noise and missing pixels. While the frozen dynamics model can handle masked inputs, the downstream linear probe (calibrated on clean data) fails when the encoded state shifts off the calibration manifold due to noise.
- The success of KE as a control objective is attributed to the specific system constraints (single scalar boundary actuator) which make the map from state to KE injective; this may not generalize to multi-actuator or topologically distinct flows without further observables.

Significance and Claims

The paper claims that latent dynamics can remain both dynamic and goal-agnostic, serving as a powerful substrate for control. However, the authors argue that calibrated observables (provided they guarantee unique continuation) are a superior objective for state control compared to raw latent Euclidean distance.

The primary significance lies in demonstrating that a world model trained without reward supervision or state reconstruction can support multiple, diverse control objectives. By separating the learned dynamics from the task metric, the framework avoids the "retraining penalty" of traditional surrogate models. The authors modestly conclude that while their specific success with Kinetic Energy is tied to the injectivity of the observable in their specific setup, the general approach of using a calibrated physical probe downstream of a reconstruction-free JEPA offers a promising path for overcoming the geometric distortions of latent spaces in PDE control.