Technische Samenvatting: SAGE – Subgoal-Conditioned Action Generation voor Latent World Model Planning

Probleemstelling

Latente wereldmodellen (latent world models) hebben zich gevestigd als efficiënte planningsparadigma's door het aanleren van actie-geconditioneerde voorspellende dynamiek in een compacte representatieruimte, waardoor agenten kandidaat-actiesequenties kunnen simuleren en evalueren zonder pixel-niveau reconstructie. Echter, naarmate de planningshorizon zich uitbreidt, wordt de effectiviteit van deze modellen steeds meer beperkt door de kwaliteit van de actievoorstellen. Onder een vaststaand budget voor kandidaten faalt de initialisatie van actiezoekopdrachten met een generieke willekeurige distributie, omdat dit onvoldoende dekking biedt voor de exponentieel expanderende actieruimte. Bijgevolg worden doelgerichte sequenties zelden gesampled, en wordt het wereldmodel niet blootgesteld aan hoogwaardige kandidaat-toekomsten voor evaluatie. Deze "proposal gap" (voorstelsgat) beperkt de prestaties bij langetermijnplanning (long-horizon planning) ernstig, zelfs wanneer de onderliggende voorspellende dynamiek accuraat is.

Methodologie: SAGE-framework

Om de proposal gap aan te pakken, introduceren de auteurs SAGE (Subgoal-Conditioned Action Generation), een prior-geconditioneerd planningsframework. SAGE vervangt de willekeurige voorstel-initialisatie door gestructureerde sturing afgeleid van voorspelde latente subdoelen (subgoals). Het framework werkt door langetermijnplanning te deconstrueren in twee nauw gekoppelde, geleerde beslissingen, terwijl het onderliggende wereldmodel (specifiek LeWM) bevroren blijft.

1. Goal-Conditioned Subgoal Generator

Het eerste component is een variabele-lengte latente subgoal generator. Gegeven de huidige observatiegeschiedenis, een ver doel en een gevraagde toekomstige duur ($\tau$), voorspelt deze module een bereikbare intermediaire latente staat (subgoal) op tijdstip $t+\tau$.

• Architectuur: Een vierlaagse Transformer-decoder die geschiedenis-tokens, de huidige staat, verre doel-tokens en temporele embeddings ($\delta, \tau$) als input neemt.
• Mechanisme: Het produceert een residuele correctie op de verre doel-latent, wat een lokale target $\hat{z}_{t+\tau}$ oplevert die bereikbaar is binnen de specifieke duur $\tau$.
• Functie: Deze generator deconstrueert verre doelen in bereikbare intermediaire staten op verschillende temporele schalen, waarbij zowel fijnmazige lokale dynamiek (korte $\tau$) als hogere niveaus van taakprogressie (lange $\tau$) worden gevangen.

2. Subgoal-Conditioned Action Generator

Het tweede component brengt de huidige geschiedenis, het verre doel en de voorspelde lokale subgoal in kaart naar een distributie van actie-opties.

• Architectuur: Een drie-laagse Transformer-decoder die geleerde actie-queries gebruikt om te attentionen op de conditionerende geheugenbron (geschiedenis, staat, subgoal, ver doel).
• Output: Het netwerk parametriseert een traject-niveau Gaussian Mixture Model (GMM), $q_\phi(a_{t:t+\tau-1} | \dots) = \sum \pi_m \mathcal{N}(a; \mu_m, \Sigma_m)$.
• Functie: In plaats van te samplen uit een generieke distributie, samplet de planner actiesequenties geconditioneerd op de specifieke lokale subgoal. Dit zorgt ervoor dat de voorgestelde acties structureel zijn afgestemd op de onmiddellijke, bereikbare target.

3. Planning en Verfijningslus

Het planningsproces integreert deze geleerde componenten met het bevroren wereldmodel:

1. Decompositie: Een langetermijn-query $H$ wordt gedeconstrueerd in een schema van duurders $(\tau_1, \dots, \tau_R)$.
2. Subgoal Predictie: In elke fase voorspelt de subgoal generator een lokale target $\hat{z}_{t+\tau}$ gebaseerd op de resterende horizon.
3. Proposal Generatie: De actiegenerator produceert $K$ kandidaat-actiesequenties geconditioneerd op deze lokale target.
4. Evaluatie & Verfijning: De bevroren LeWM voorspelt de latente toekomsten van deze kandidaten en evalueert hun kosten ten opzichte van de lokale subgoal. De Cross-Entropy Method (CEM) verfijnt vervolgens de elite-set van kandidaten over 30 rondes.
5. Executie: De agent voert het leidende segment van de geselecteerde sequentie uit, observeert de nieuwe staat en plant opnieuw.

Belangrijkste Bijdragen

1. Multi-Horizon Latent Subgoal Generator: Een nieuw module dat verre doelen deconstrueert in bereikbare intermediaire staten op verschillende temporele schalen, waardoor de kloof tussen de huidige staat en verre doelstellingen effectief wordt overbrugd.
2. Subgoal-Conditioned Action Generator: Een mechanisme dat actiekandidaten genereert geconditioneerd op voorspelde subgoals, waarbij generieke willekeurige initialisatie wordt vervangen door gestructureerde, taakbewuste priors.
3. Empirische Validatie van Gestructureerde Voorstellen: Demonstratie dat het koppelen van variabele-lengte subgoal generatie met geconditioneerde actiegeneratie de langetermijnplanning aanzienlijk verbetert, terwijl de kortetermijnprestaties behouden blijven, zonder de voorspellende dynamiek van het bevroren wereldmodel te wijzigen.

Experimentele Resultaten

De methode werd geëvalueerd op PushT (vereist aanhoudend vlak contact) en OGBench Cube (vereist armbeweging, grijpen en plaatsen) met een gecontroleerde setup met een bevroren LeWM-checkpoint, identieke kandidaat-budgetten (300) en CEM-verfijning (30 rondes).

• Langetermijnprestaties: Naarmate de planningshorizon toeneemt, vertoont SAGE substantiële winst ten opzichte van de baselines (standaard LeWM en PRISM).
  • Bij een target offset van H=150:
    • PushT: Succespercentage verbeterd van 12,7% (LeWM) naar 64,7% (SAGE).
    • OGBench Cube: Succespercentage verbeterd van 26,7% (LeWM) naar 67,3% (SAGE).
• Kortetermijnprestaties: SAGE behoudt sterke prestaties op korte horizonten (bijv. 94,0% op PushT bij H=25), wat aangeeft dat de gestructureerde priors de prestaties niet verslechteren wanneer het doel nabij is.
• Ablatie-studies:
  • Generator-only: Het gebruik van gegenereerde lokale targets met een standaard Gaussische voorstel (zonder de geleerde actiegenerator) leverde nog steeds significante verbeteringen op (bijv. 58,7% op PushT bij H=150 versus 12,7%), wat de waarde van lokale target-decompositie bevestigt.
  • SAGE Prior Top: Het uitvoeren van de optie met het hoogste gewicht uit de geleerde prior zonder LeWM-ranking of CEM-verfijning resulteerde in slechte langetermijnprestaties (16,0% bij H=150), wat het belang van wereldmodel-gestuurde zoekopdrachten en verfijning benadrukt.
• Temporele Schema's: Experimenten met verschillende duur-schema's (bijv. variërend in de mix van 15, 20 en 25-staps commitments) toonden aan dat zowel de frequentie van herplannen als de ordening van de duurders een significante impact heeft op de succespercentages.

Betekenis en Claims

Het artikel claimt dat SAGE een praktisch pad biedt om de planningsmogelijkheden van bestaande latente wereldmodellen uit te breiden zonder de computationele kosten van het hertrainen van hun voorspellende dynamiek. Door de structuur van de actievoorstellen te verbeteren — specifiek door verre doelen te deconstrueren in bereikbare lokale targets en de actiegeneratie te conditioneren op deze targets — mitigeert het framework effectief de bottleneck in de kwaliteit van de voorstellen die inherent is aan langetermijnplanning.

De auteurs benadrukken dat hun aanpak de "voorstel"-mechanisme scheidt van de "voorspellings"-mechanisme. Terwijl het wereldmodel een vaste simulator blijft voor het evalueren van toekomsten, zorgt SAGE ervoor dat de kandidaten die aan de simulator worden gepresenteerd plausibel zijn en in lijn zijn met de taakstructuur. Deze modulaire verbetering maakt aanzienlijke winst mogelijk in complexe, langdurige robotica-taken, terwijl de efficiëntie en robuustheid van het onderliggende wereldmodel behouden blijven.