Imagine-then-Plan: Agent Learning from Adaptive Lookahead with World Models

Titel: Imagine-then-Plan (ITP): Agent Learning from Adaptive Lookahead with World Models

Auteurs: Youwei Liu, Jian Wang, et al. (The Hong Kong Polytechnic University & Central South University)

1. Het Probleem

Huidige autonome agenten, vaak gebaseerd op Large Language Models (LLMs), kampen met een tekort aan "diepe verankering" (shallow grounding). Hoewel ze goed kunnen redeneren op basis van directe observaties en historische interacties, ontbreekt het hen vaak aan een causaal begrip van hoe hun huidige acties de toekomstige staat van het milieu zullen beïnvloeden.

• Beperkingen van bestaande methoden: Bestaande wereldmodellen (world models) worden vaak gebruikt voor vaste, korte roll-outs (enkelvoudige stappen of een vast aantal stappen). Dit is suboptimaal voor complexe taken met lange horizonnen (zoals huishoudelijke taken of wetenschappelijke experimenten), omdat ze lange-termijn afhankelijkheden missen.
• Efficiëntieprobleem: Het uitvoeren van diepe roll-outs voor elke stap is computationally duur en onnodig voor triviale acties.
• Gevolg: Zonder het vermogen om de toekomst te projecteren, lopen agenten het risico op catastrofale fouten die pas worden ontdekt nadat acties onomkeerbaar zijn uitgevoerd.

2. Methodologie: Imagine-then-Plan (ITP)

De auteurs stellen Imagine-then-Plan (ITP) voor, een unificerend kader dat agenten in staat stelt om proactief te "repeteren" in een mentale sandbox voordat ze handelen.

A. Conceptuele Innovatie: POIMDP

In plaats van het traditionele Partially Observable Markov Decision Process (POMDP), introduceren de auteurs het Partially Observable and Imaginable Markov Decision Process (POIMDP).

• In dit kader wordt het beleidsbeleid ($\pi_\theta$) geoptimaliseerd op basis van een dubbelstroomrepresentatie: de concrete, waarneembare huidige staat ($s_t$) én de vooruitziende, verbeelde toekomst ($\hat{\tau}_t$).
• De agent gebruikt een geleerd wereldmodel ($M_\phi$) om multi-stap trajecten te genereren binnen een "mentale sandbox".

B. Adaptieve Lookahead

Een kerninnovatie is de adaptieve lookahead-mechanisme. In plaats van een vast aantal stappen ($k$) te gebruiken, bepaalt de agent dynamisch hoe ver hij moet kijken ($K_t$) op basis van:

• De geschatte voortgang naar het einddoel.
• De complexiteit van de huidige beslissing.
• Een afweging tussen het bereiken van het doel en de rekenkosten.

C. Twee Varianten van ITP

Het kader wordt geïmplementeerd in twee varianten:

1. ITP-Inf (In-Imagination Learning, Training-vrij):

   • Zowel het beleidsmodel als het wereldmodel zijn bevroren (geen extra training).
   • Tijdens inferentie voert de agent een drie-staps procedure uit:
     1. Adaptieve horizonselectie: De agent kiest een $K_t$ gebaseerd op de huidige context.
     2. Wereldmodel-imaginatie: Het genereren van een toekomsttraject $\hat{\tau}_t$ via het wereldmodel.
     3. Reflectieve beleidsgeneratie: De agent reflecteert op dit verbeelde traject om conflicten of bottlenecks te detecteren en past zijn actie aan voordat hij in de echte omgeving handelt.

2. ITP-RL (Reinforced Training):

   • Een variant die het agentbeleid en de lookahead-selectie gezamenlijk optimaliseert via versterking leren.
   • Fase 1 (Pseudo-labeling): Het genereren van "juiste" lookahead-horizons voor experttrajecten door verschillende roll-outs te scoren.
   • Fase 2 (Warm-up): Supervised fine-tuning van het beleidsmodel en een lichte "K-head predictor" (een lineaire laag die de optimale $K_t$ voorspelt).
   • Fase 3 (Online Optimalisatie): Gebruik van het Advantage Actor-Critic (A2C) algoritme. De beloning functie straft zowel het falen van de taak als de rekenkosten (token-gebruik) af, waardoor de agent leert wanneer het de moeite waard is om diep na te denken.

3. Belangrijkste Bijdragen

1. Conceptualisering van POIMDP: Een nieuwe theoretische basis die verbeelde toekomstige trajecten integreert in de besluitvorming van agenten.
2. ITP Framework: Een unificerend systeem dat wereldmodel-imaginatie combineert met een adaptieve lookahead-mechanisme voor deliberatieve planning.
3. Empirische Validatie: Bewijs dat zowel de training-vrije als de getrainde varianten de succespercentages van LLM-agenten aanzienlijk verbeteren op complexe, lange-horizon taken.
4. Open Source: Publicatie van code en data om reproduceerbaarheid te garanderen.

4. Resultaten

De auteurs hebben ITP getest op vier representatieve benchmarks: ALFWorld (embodied household tasks), ScienceWorld (wetenschappelijke experimenten), WebShop (webnavigatie) en StableToolBench (tool use).

• Prestaties:
  • ITP-Inf (Training-vrij): Overtreft sterke prompting-baselines (zoals ReAct en RAP) aanzienlijk zonder extra training. Bijvoorbeeld, op ALFWorld steeg het succespercentage van ~28% (RAP) naar ~36% (ITP-Inf) met Qwen2.5-7B.
  • ITP-RL (Getraind): Bereikt de hoogste succespercentages in alle geteste backbone-modellen. Met Qwen3-8B op ALFWorld werd een succespercentage van 88,57% bereikt, wat een verbetering is ten opzichte van bestaande trainingsmethoden zoals SFT en IWM.
• Efficiëntie vs. Effectiviteit:
  • Ablatiestudies: Het verwijderen van de online versterkte training (RT) leidt tot een drastische daling in prestaties (van 88,57% naar 71,42% op ALFWorld), wat aantoont dat het leren wanneer te imagineren cruciaal is.
  • Vaste vs. Adaptieve Lookahead: Een vaste lookahead ($k$) is breekbaar; te weinig stappen missen lange-termijn afhankelijkheden, te veel stappen verhogen de kosten en versterken modelfouten. De adaptieve lookahead van ITP bereikt een hogere succesrate met een aanzienlijk lagere rekenkosten (gemeten in "Normalized Budget").
  • Robuustheid: ITP presteert consistent goed over verschillende backbone-modellen (Qwen, Llama, DeepSeek), hoewel de training-vrije variant meer afhankelijk is van de kwaliteit van het wereldmodel.

5. Betekenis en Conclusie

Dit onderzoek markeert een verschuiving van reactieve agenten naar deliberatieve agenten.

• Sluiting van de Verankering: Door agenten toe te staan om toekomstige gevolgen te simuleren voordat ze handelen, wordt het probleem van "shallow grounding" aangepakt.
• Efficiënte Planning: De adaptieve lookahead zorgt ervoor dat rekenkracht alleen wordt ingezet wanneer het nodig is voor complexe beslissingen, wat de schaalbaarheid voor real-world toepassingen verbetert.
• Toekomstperspectief: Hoewel de huidige evaluatie zich richt op tekstuele omgevingen, biedt ITP een fundamenteel kader dat kan worden uitgebreid naar multimodale en robotische omgevingen. De auteurs erkennen echter dat de extra inferentie-tijd een uitdaging blijft voor real-time toepassingen.

Kortom, Imagine-then-Plan biedt een principieel geavanceerde aanpak voor het gebruik van wereldmodellen, waarbij agenten leren om hun verbeelding strategisch in te zetten om complexe taken succesvoller en efficiënter te voltooien.