Technische Samenvatting: FF-JEPA: Lange-termijn Planning in World Models met Latente Planners

Probleemstelling

Recente vooruitgang in visuele wereldmodellen, met name Joint Embedding Predictive Architectures (JEPAs) zoals LeWorldModel (LeWM), hebben het vermogen aangetoond om toekomstige observaties te voorspellen en modelgebaseerde controle te ondersteunen via sampling-gebaseerde planners zoals de Cross-Entropy Method (CEM). De auteurs identificeren echter drie kritieke beperkingen die deze "platte" wereldmodellen verhinderen bij de inzet voor real-world, lange-termijn taken:

1. Computationele Prohibitieve Factoren: Plannen over lange horizonten vereist herhaalde rollouts en steeds grotere actiesequenties, wat CEM-optimalisatie computationeel onhandelbaar maakt.
2. Foutaccumulatie (Error Compounding): Autoregressieve voorspelfouten accumuleren over vele stappen, waardoor de planner divergeert bij complexe trajecten.
3. Doelafhankelijkheid: Bestaande methoden vereisen doorgaans een expliciete doelafbeelding of -toestand ($o_g$) om planning aan te sturen. In veel real-world scenario's is een dergelijke concrete doelspecificatie niet beschikbaar of onpraktisch.

Hoewel hiërarchische methoden en generatieve benaderingen hebben geprobeerd deze problemen aan te pakken, opereren zij vaak nog steeds in de pixelruimte of vertrouwen zij op extern gespecificeerde doelafbeeldingen/trajecten, wat de toepasbaarheid in open-ended omgevingen beperkt.

Methodologie: Forward-Forward JEPA (FF-JEPA)

Het paper stelt FF-JEPA voor, een hiërarchisch framework dat JEPA-stijl wereldmodellen uitbreidt door een latente planner ($G$) te introduceren die naast een standaard actie-geconditioneerd forward model ($P$) opereert. Deze aanpak deelt lange-termijn planning op in een sequentie van behapbare, kortetermijn optimalisatieproblemen zonder dat daar een expliciete doelafbeelding voor nodig is.

Kernarchitectuur

Het systeem bestaat uit twee primaire componenten die getraind worden binnen de latente ruimte gedefinieerd door een bevroren, vooraf getrainde wereldmodel-encoder ($E$):

1. Het Wereldmodel ($P$): Een standaard actie-geconditioneerd forward dynamics model (gebaseerd op LeWM) dat de volgende latente toestand voorspelt gegeven een historie van toestanden en een actie: $\hat{z}_{t+1} = P(z_{t-W_P+1:t}, a_t)$.
2. De Latente Planner ($G$): Een actie-vrij forward dynamics model dat de volgende subgoal latent ($\hat{z}_{sg}$) voorspelt op basis van enkel een historie van voorgaande subgoal-latenten. Het opereert op een hogere temporele resolutie (elke $H$ omgevingstappen).
   • Deterministische Planner ($G_{Det}$): Een transformer-architectuur (vergelijkbaar met de predictor van het wereldmodel maar zonder actie-conditionering) getraind om de gemiddelde kwadratische fout tussen voorspelde en ware toekomstige subgoals te minimaliseren.
   • Diffusion Planner ($G_{DM}$): Een DiT (Diffusion Transformer) backbone getraind met een denoising score-matching objectief om een sequentie van toekomstige subgoals te voorspellen.

Planning Inferentie

Het inferentieproces is hiërarchisch:

1. Subgoal Voorspelling: De latente planner $G$ voorspelt de volgende subgoal-toestand $\hat{z}_{sg, m+1}$ gegeven de huidige subgoal-historie.
2. Actie Optimalisatie: Het wereldmodel $P$ gebruikt CEM om een actiesequentie ($a_{1:H}$) te vinden die de voorspelde toestand naar de voorspelde subgoal $\hat{z}_{sg, m+1}$ stuurt binnen een kort horizon $H$.
3. Iteratie: Deze cyclus herhaalt zich, waardoor korte-termijn plannen effectief worden geketend om lange-termijn doelstellingen te bereiken.

Cruciaal is dat deze methode de noodzaak voor een doelafbeelding tijdens inferentie wegneemt; de planner genereert het traject naar een impliciet gedefinieerd doel.

Belangrijkste Bijdragen

• Doel-vrije Lange-termijn Planning: Het framework maakt planning in wereldmodellen mogelijk zonder expliciete doelafbeeldingen door een latente planner te gebruiken om tussenliggende subgoals te genereren.
• Hiërarchische Decompositie: Het pakt de "long-horizon collapse" van platte wereldmodellen aan door complexe trajecten op te splitsen in korte, oplosbare subproblemen, wat foutaccumulatie en computationele kosten mitigeert.
• Ongelabelde Training: De latente planner kan getraind worden op ongelabelde succesvolle demonstraties (gesubsampleerd uit de latente ruimte van het wereldmodel), wat de afhankelijkheid van actie-gelabelde data vermindert.
• Verenigd Framework: Het herinterpreteert het wereldmodel als een inverse dynamics module die opereert over geïmagineerde latente trajecten, waardoor predictieve modellering en controle verenigd worden binnen één coherent latente ruimte.

Empirische Evaluatie

De auteurs hebben FF-JEPA geëvalueerd op de PushT taak met drie experimentele opstellingen: Korte-termijn (25 stappen), Lange-termijn (75 stappen) en Random Initialisatie (300 stappen met gerandomiseerde startposities).

Resultaten

• Lange-termijn Prestaties: De baseline flat LeWM stortte in bij lange-termijn taken (3,52% succesrate bij $t=75$) en faalde volledig onder random initialisatie (0,00%). In tegenstelling hiertoe bereikte FF-JEPA (Diffusion) 91,80% succes bij $t=75$ en 82,42% onder random initialisatie.
• Voordeel van Doel-vrijheid: In tegenstelling tot baselines (LeWM, DINO-WM) die een doelafbeelding vereisen voor inferentie, heeft FF-JEPA dat niet nodig. Ondanks dit presteerde het beter dan hiërarchische DINO-WM baselines in zowel korte-termijn (96,09% vs. 89,0%) als lange-termijn (91,80% vs. 61,0%) settings.
• Data Efficiëntie: Ablatie-studies toonden aan dat het trainen van de diffusion planner op slechts 200 hoogwaardige demonstraties (een 40x reductie in data) nog steeds competitieve prestaties opleverde (76,17% vs. 82,42% met volledige data), wat wijst op een hoge data-efficiëntie.
• Computationele Overhead: De deterministische planner voegde verwaarloosbare inferentietijd toe (2,1ms) vergeleken met CEM (926ms), terwijl de diffusion planner ~242ms toevoegde. Beiden zijn significant sneller dan de herhaalde rollouts die vereist zijn voor lange-termijn flat planning.

Foutanalyse

Falen trad meestal op wanneer de agent te ver afdrifte van de voorspelde subgoal positie of wanneer de latente planner een kwalitatief minderwaardige subgoal-toestand genereerde (bijv. de agent missend).

Betekenis en Claims

Het paper positioneert FF-JEPA als een veelbelovende richting voor schaalbare, doel-vrije robot learning. Door aan te tonen dat actie-geconditioneerde wereldmodellen gekoppeld aan latente planners getraind op ongelabelde data de beperkingen van flat planning kunnen overwinnen, suggereren de auteurs dat deze aanpak agents in staat stelt om te opereren in open-ended omgevingen waar expliciete doelspecificaties niet beschikbaar zijn. Het werk benadrukt dat het opdelen van planning in kortetermijn subgoals een levensvatbare strategie is om de computationele en fout-accumulerende problemen van huidige wereldmodel-benaderingen te mitigeren.