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

Technische Samenvatting: Naar doel-agnostische Joint-Embedding Predictive Control van Partiële Differentiaalvergelijkingen

Probleemstelling

Het aansturen van systemen die worden beheerst door partiële differentiaalvergelijkingen (PDE's), zoals vloeistofdynamica, is uitdagend vanwege hoogdimensionele toestandsruimten, niet-lineaire relaties en de computationele kosten van numerieke simulaties die nodig zijn voor het evalueren van controle-sequenties. Bestaande methoden kampen met specifieke beperkingen:

• Model-gebaseerde controle vereist vaak probleem-specifieke analyse of gereduceerde-orde modellen (ROM's) die mogelijk een gebrek aan algemeenheid vertonen.
• Geleerde surrogaatmodellen worden doorgaans getraind op voorwaartse simulatie in de toestandsruimte, geoptimaliseerd voor specifieke doelen, wat hertraining noodzakelijk maakt wanneer doelen veranderen en de moeilijkheid van volledige toestandreconstructie met zich meebrengt.
• Reinforcement Learning (RL) presteert goed op gedefinieerde taken, maar heeft moeite met generalisatie door de gevoeligheid voor de selectie van de beloningsfunctie.

De auteurs stellen een raamwerk voor dat het leren van systeemdynamica loskoppelt van het controledoel, met als doel een "doel-agnostische" controller die kan adapteren aan nieuwe taken zonder de onderliggende wereldmodellen te hoeven hertrainen.

Methodologie

Het artikel introduceert een controleframework gebouwd rond een Joint-Embedding Predictive Architecture (JEPA) toegepast op PDE's. De kernfilosofie is het leren van een latent wereldmodel dat voorspelt hoe acties tussen latente toestanden in kaart worden gebracht, zonder de fysieke toestand te reconstrueren of beloningssignalen te gebruiken tijdens de training.

1. Doel-agnostische training van het wereldmodel:

• Architectuur: Een lichtgewicht 2D Vision Transformer (ViT) encoder ($E_\theta$) en een actie-geconditioneerde latente predictor ($g_\phi$) worden offline getraind op toestands-transities.
• Trainingsdoel: Het model wordt getraind met gemaskeerde observaties (het willekeurig maskeren van 30–90% van het veld) om toekomstige latente toestanden te voorspellen. De verliesfunctie bevat:
  • Rollout Loss: Het minimaliseren van de afstand tussen voorspelde en werkelijke toekomstige latenten.
  • VICReg: Variantie- en covariantie-straffen om representatie-instorting te voorkomen.
  • Temporal Straightening: Het stimuleren van latente trajecten om direct te zijn in plaats van omwegen te maken.
  • Actie-gevoeligheid: Een hulp-decoder voorspelt acties vanuit latente verschuivingen (geïnspireerd door Delta-JEPA) om te waarborgen dat de latente ruimte gevoelig is voor controle-inputs.
• Belangrijke beperking: Er is geen decoder die latente toestanden terug naar het fysieke veld $U$ brengt, en er wordt geen beloningsfunctie gebruikt tijdens de training. De encoder en de predictor worden bevroren na de training.

2. Implementatie en Controlestrategie:

• Model-Predictive Path Integral (MPPI): Een bevroren predictor wordt gebruikt binnen een MPPI-controller om kandidaat actie-sequenties te evalueren.
• Doel-interfaces: Het framework ondersteunt drie implementatiemodi zonder de gewichten van het wereldmodel te wijzigen:
  1. Tijdsafhankelijke tracking: Het volgen van een sequentie van gecodeerde toestanden of fysieke observaties.
  2. Vaste-doelstabilisatie: Het handhaven van een specifieke toestandconfiguratie.
  3. Regulatie van kinetische energie (KE): Het stabiliseren rond een specifiek energieniveau.
• De Observable-Aligned Controller: Een cruciale innovatie is het gebruik van een lichtgewicht lineaire probe ($q_\omega$) die getraind is op bevroren latente rollouts om een fysieke observeerbare grootheid (specifiek: kinetische energie) te schatten.
  • In plaats van de ruwe Euclidische afstand ($L_2$) in de latente ruimte te minimaliseren (wat mogelijk niet overeenkomt met de fysieke fout), minimaliseert de MPPI-kostenfunctie de fout tussen de voorspelde fysieke observeerbare en de doel-observeerbare.
  • Deze aanpak berust op de aanname dat de gekozen observeerbare (KE) garandeert dat er een unieke voortzetting (unique continuation) bestaat voor de specifieke systeemrestricties.

Belangrijkste Bijdragen

1. Eerste JEPA voor gesloten-lus PDE-controle: De auteurs presenteren de eerste implementatie van gesloten-lus PDE-controle met een joint-embedding latent wereldmodel dat actie-geconditioneerde dynamica leert, onafhankelijk van de downstream-taak.
2. Observable-Aligned Planning: Zij introduceren een methode waarbij een goedkope lineaire probe een fysieke scalaire grootheid (kinetische energie) schat uit latente voorspellingen, en de controller de trackingfout in fysieke eenheden optimaliseert in plaats van de latente afstand.
3. Doel-agnostiek via Bevroren Modellen: Een enkel bevroren model checkpoint en een korte latente horizon van 5 stappen kunnen succesvol tijdsafhankelijke tracking, vaste-veld stabilisatie en kinetische energie-regulatie afhandelen zonder taak-specifieke verfijning of hertraining.
4. Aantoonbare Prestatiewinsten: Op de PDE Control Gym 2D Navier–Stokes benchmark presteert de KE-probe aanpak beter dan de latente-$L_2$ planning, waarbij de natuurlijke beloning wordt verbeterd en de veld-RMSE wordt verminderd.

Experimentele Resultaten

Het framework werd geëvalueerd op de PDE Control Gym 2D Navier–Stokes omgeving (eenheid vierkant, $21 \times 21$ rooster).

• Benchmark Prestaties (50 episodes):
  • De overstap van de latente-$L_2$ kosten naar de KE-probe kosten verbeterde de gemiddelde natuurlijke beloning van −12,08 naar −10,90.
  • De laatst-kwartiel snelheidveld-RMSE daalde van 0,0765 naar 0,0692.
  • De KE-probe controller won 49 van de 50 gepaarde episodes tegen de latente-$L_2$ baseline.
• Aperiodische Doelen:
  • Getest tegen drie verschillende, niet-geziene, aperiodische doelsignalen.
  • KE-planning verminderde de late-veld RMSE met 53% (0,0220 vs. 0,0469) vergeleken met latente-$L_2$ planning, en won alle 30 gepaarde episodes.
  • Het corrigeerde succesvol de faseverschuivingen (phase lags) die werden waargenomen bij de latente-$L_2$ aanpak.
• Stabilisatie:
  • Voor vaste-veld stabilisatie bereikte de KE-probe een gemiddelde relatieve fout van 2,7% bij het besturen van het systeem naar een stabiele configuratie, waarbij het evenwichtsevenwicht met hoge precisie werd hersteld.
• Beperkingen:
  • De KE-probe is bros tegen meetruis en ontbrekende pixels. Hoewel het bevroren dynamica-model met gemaskeerde inputs kan omgaan, faalt de downstream lineaire probe (gekalibreerd op schone data) wanneer de gecodeerde toestand van de kalibratie-manifold afwijkt door ruis.
  • Het succes van KE als controoldoel wordt toegeschreven aan de specifieke systeemrestricties (enkele scalaire randactuator) die de kaart van toestand naar KE injectief maken; dit kan mogelijk niet generaliseren naar systemen met meerdere actuatoren of topologisch andere stromingen zonder verdere observeerbare grootheden.

Betekenis en Claims

Het artikel claimt dat latente dynamica zowel dynamisch als doel-agnostisch kan blijven, en dient als een krachtig substraat voor controle. De auteurs stellen echter dat gekalibreerde observeerbare grootheden (mits zij unieke voortzetting garanderen) een superieur doel zijn voor toestandcontrole vergeleken met de ruwe latente Euclidische afstand.

De primaire betekenis ligt in het bewijs dat een wereldmodel dat getraind is zonder beloningssupervisie of toestandreconstructie, diverse en uiteenlopende controledoelen kan ondersteunen. Door de geleerde dynamica te scheiden van de taakmetriek, vermijdt het framework de "hertrainingsstraf" van traditionele surrogaatmodellen. De auteurs concluderen bescheiden dat hoewel hun specifieke succes met kinetische energie verbonden is aan de injectiviteit van de observeerbare in hun specifieke opstelling, de algemene aanpak van het gebruik van een gekalibreerde fysieke probe achter een reconstructie-vrije JEPA een veelbelovende weg biedt om de geometrische vervormingen van latente ruimten in PDE-controle te overwinnen.