Technische Samenvatting: AdaJEPA: Een Adaptief Latente Wereldmodel

Probleemstelling

Latente wereldmodellen, in het bijzonder Joint-Embedding Predictive Architectures (JEPA's), maken efficiënte planning mogelijk door toekomstige toestanden te voorspellen in een compacte latente ruimte. Deze modellen worden doorgaans getraind op beloningsvrije offline trajecten en gekoppeld aan Model Predictive Control (MPC) voor doelgeconditioneerde taken. Er bestaat echter een kritieke beperking: het wereldmodel is meestal bevroren na de training.

Wanneer het model wordt ingezet, kunnen de voorspellingen in nauwkeurigheid afnemen door distributieverschuivingen tijdens de testfase (bijv. visuele veranderingen zoals ruis of verlichting, of dynamische veranderingen zoals wrijving en massa), waardoor de MPC-planner een incorrect doel optimaliseert. Acties die effectief lijken in de "latente verbeelding", kunnen falen in de echte omgeving. Kleine fouten in één stap stapelen zich op over de planningshorizon, wat leidt tot aanzienlijke prestatievermindering. Huidige benaderingen vereisen vaak aanvullende expertdemonstraties, fine-tuning data voor de doel-domein, of complexe outer loops om dit aan te pakken, wat de toepassing in de echte wereld belemmert.

Methodologie: AdaJEPA

De auteurs stellen AdaJEPA voor, een adaptief latent wereldmodel dat test-time adaptatie uitvoert binnen de gesloten lus van MPC. In plaats van het model als statisch te behandelen, kalibreert AdaJEPA het wereldmodel continu met behulp van zelfgesuperviseerde signalen die voortkomen uit de eigen interacties van de agent.

De Plan–Execute–Adapt–Replan Lus

Het kernmechanisme werkt als volgt bij elke MPC-stap:

1. Plan: De agent plant een reeks acties met behulp van het huidige (mogelijk verouderde) wereldmodel om een latente doelbereikingskosten te minimaliseren.
2. Execute: De agent voert de eerste actiechunk ($a_t$) uit in de omgeving.
3. Observe: De agent ontvangt de volgende observatie ($o_{t+1}$).
4. Adapt: De transitie $\{o_t, a_t, o_{t+1}\}$ wordt toegevoegd aan een replay buffer. De modelparameters worden bijgewerkt via een klein aantal gradiëntstappen (vaak slechts één) om de latente voorspellingsfout op deze nieuwe transitie te minimaliseren.
5. Replan: Het bijgewerkte model wordt onmiddellijk gebruikt voor de volgende planningsiteratie.

Technische Componenten

• Architectuur: Het onderliggende wereldmodel bestaat uit een sensorische encoder ($\mathcal{E}_s$), een actie-encoder ($\mathcal{E}_a$) en een voorspeller ($f_\theta$). Het wordt offline getraind om toekomstige latente representaties ($z_{t+1}$) te voorspellen in plaats van pixels te reconstrueren.
• Adaptatieverlies: De adaptatie gebruikt hetzelfde zelfgesuperviseerde latente voorspellingsdoel als de pretraining:
  $$\mathcal{L}_{ada}(\mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{(o_i, a_i, o_{i+1}) \in \mathcal{B}} \ell(f_\theta(z_i, \mathcal{E}_a(a_i)), \text{sg}(z_{i+1}))$$
  waarbij $\mathcal{B}$ een buffer is van recente transities, $\ell$ een latente voorspellingsverlies is (bijv. MSE), en $\text{sg}(\cdot)$ een stop-gradient operator is om instorting te voorkomen.
• Parameterupdates: Om efficiëntie te garanderen, is de adaptatie lichtgewicht. De auteurs updaten alleen een kleine subset van de parameters (bijv. de laatste lagen van de encoder en de voorspeller) in plaats van het gehele netwerk.
• Bufferstrategieën: De replay buffer $\mathcal{B}$ kan worden beheerd via "recent-N" (het bijhouden van de meest recente $N$ transities) of "hard-N" (het bijhouden van transities met de grootste voorspellingsfouten).

Belangrijkste Bijdragen

1. Eerste JEPA-adaptatie tijdens planning: Naar weten van de auteurs is dit het eerste werk dat een JEPA-wereldmodel aanpast tijdens de planningsfase, in plaats van uitsluitend te vertrouwen op pre-training of offline fine-tuning.
2. Closed-Loop Zelf-supervisie: De methode maakt gebruik van de directe gevolgen van uitgevoerde acties als een zelfgesuperviseerd signaal, waardoor de noodzaak voor beloningslabels, expertdemonstraties of aparte dataverzamelfases vervalt.
3. Sample- en Computerefficiëntie: Het framework bereikt significante prestatiewinsten met minimale computationele overhead. In veel experimenten is een enkele gradiëntstap per MPC-replan stap voldoende.
4. Robuustheid tegen Distributieverschuivingen: De benadering demonstreert robuustheid tegen visuele verschuivingen (ruis, onscherpte, kleurveranderingen), dynamische verschuivingen (massa, demping) en structurele verschuivingen (ongeziene doolhoflay-outs).

Experimentele Resultaten

De auteurs evalueerden AdaJEPA op de PushT (blok duwen) en PointMaze (navigatie) benchmarks onder diverse train-test distributieverschuivingen.

• In-Distribution Prestaties: Zelfs op taken die vergelijkbaar zijn met de trainingsdata, verbetert AdaJEPA de succesratio's (bijv. >20% winst op specifieke PushObj vormen) of behoudt het de baseline-prestaties zonder degradatie, wat bewijst dat het veilig is om in-distribution toe te passen.
• Out-of-Distribution (OOD) Prestaties:
  • Vormverschuivingen: Bij ongeziene blokvormen verdubbelde AdaJEPA de planningssuccesratio bijna vergeleken met het bevroren model.
  • Visuele Verschuivingen: De methode toonde duidelijke winsten onder Gaussische onscherpte, salt-and-pepper ruis en donkere verlichting.
  • Dynamiek & Layout Verschuivingen: In PointMaze herstelde AdaJEPA van fouten veroorzaakt door lage massa of hoge demping en verbeterde het de succesratio's aanzienlijk op ongeziene doolhoflay-outs (bijv. +25,3% winst op specifieke layoutverschuivingen).
• Latentie: De adaptatiestap introduceert verwaarloosbare latentie (ongeveer 0,01–0,03s per replan) omdat het slechts een kleine subset van de parameters bijwerkt. Bovendien, door sneller te herstellen van voorspellingsfouten, bereikt de agent vaak het doel in minder totale replans, wat de totale executietijd vermindert.
• Generalisatie naar Datatekort: In scenario's met weinig data compenseerde test-time adaptatie voor de beperkte trainingsdekking. Bijvoorbeeld, met beperkte trainingsdata overtrof AdaJEPA een bevroren model dat getraind was met 16x meer trajecten.

Betekenis en Claims

Het artikel betoogt dat latente wereldmodellen niet vast moeten blijven staan na de training, maar continu moeten verbeteren van de ervaring die tijdens de inzet wordt opgedaan.

• Biologische Inspiratie: De benadering is geworteld in biologische sensorimotorische adaptatie, waarbij systemen (zoals de cerebellum) hun gedrag aanpassen op basis van veranderende inputs en dynamiek.
• Complement aan Data Scaling: Hoewel grotere en diversere trainingsdata de generalisatie verbeteren, biedt AdaJEPA een sample-efficiënte route om modellen tijdens de inzet te verfijnen, wat dient als een complement aan data scaling.
• Toekomstige Richting: De auteurs suggereren dat de volgende stap het combineren van deze lichtgewicht test-time adaptatie met continue en actieve leerprocessen is om de dekking van het wereldmodel in de loop van de tijd uit te breiden, wat leidt tot meer veerkrachtige perceptie en planning in veranderende omgevingen.

De auteurs benadrukken dat hun methode agnostisch is voor de specifieke implementatie van het wereldmodel, waarbij consistente verbeteringen worden getoond over verschillende architecturen (ResNet, Transformer, DINO-gebaseerd) en planners (Gradient Descent, CEM).