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

1. Problématique

Les agents autonomes basés sur les grands modèles de langage (LLM) actuels souffrent souvent d'un ancrage superficiel (shallow grounding). Bien qu'ils puissent raisonner sur des observations immédiates et des traces d'interaction historiques, ils manquent d'une compréhension causale profonde de la manière dont leurs actions présentes façonneront l'environnement à long terme.

Les méthodes existantes utilisant des modèles du monde (World Models) pour simuler l'avenir se heurtent à deux limitations majeures :

1. Horizons de simulation rigides : Elles utilisent généralement des déroulements (rollouts) d'un seul pas ou d'un horizon fixe. Cela est sous-optimal car il peut manquer des dépendances à long terme dans des tâches complexes ou, inversement, gaspiller des ressources computationnelles sur des actions triviales.
2. Manque d'adaptabilité : Elles ne parviennent pas à équilibrer dynamiquement le compromis entre la fiabilité de la prédiction (qui diminue avec la profondeur de l'horizon) et le coût computationnel.

L'objectif est donc de doter les agents d'une capacité de prévision adaptative, leur permettant de "répéter mentalement" (rehearse) des trajectoires futures pour prendre des décisions plus délibérées et éviter les erreurs irréversibles.

2. Méthodologie : Imagine-then-Plan (ITP)

Les auteurs proposent ITP, un cadre unifié qui permet aux agents d'apprendre via une imagination de prévision (lookahead imagination). Le cœur de la méthode repose sur trois piliers :

A. Le POIMDP (Partially Observable and Imaginable MDP)

Les auteurs étendent le cadre classique du POMDP (Processus de Décision Markovien Partiellement Observable) vers un POIMDP. Dans ce nouveau formalisme, la politique de l'agent $\pi_\theta$ ne dépend pas seulement de l'état observable actuel $s_t$, mais aussi d'une trajectoire imagée future $\hat{\tau}^{(K_t)}_t$ générée par un modèle du monde.
La décision est formulée comme :
$$a_t \sim \pi_\theta(\cdot \mid s_t, \hat{\tau}^{(K_t)}_t)$$
Cela permet à l'agent d'anticiper les progrès vers l'objectif ou de détecter des conflits potentiels avant d'agir dans le monde réel.

B. Mécanisme de Prévision Adaptative (Adaptive Lookahead)

Au lieu d'utiliser un horizon fixe $K$, ITP introduit un mécanisme qui ajuste dynamiquement le nombre d'étapes d'imagination $K_t$ en fonction de l'étape de la tâche et de l'objectif final.

• Principe : L'agent alloue une prévision profonde aux décisions critiques (à fort enjeu) et une prévision minimale (ou nulle) aux actions triviales.
• Bénéfice : Cela optimise le compromis entre la précision de la prédiction (qui se dégrade avec l'horizon) et le coût computationnel.

C. Deux Variantes d'Implémentation

L'article présente deux façons d'instancier ITP :

1. ITP-Inférence (ITP-I) - Sans entraînement :

   • Les modèles (politique et monde) sont figés.
   • À chaque étape, l'agent sélectionne un horizon $K_t$, génère une trajectoire imagée via le modèle du monde, puis utilise cette trajectoire pour une réflexion (reflection) afin d'affiner son action finale.
   • C'est une méthode "plug-and-play" pour améliorer les agents LLM existants.

2. ITP-Renforcement (ITP-R) - Entraînement par Renforcement :

   • L'agent apprend explicitement quand et combien de temps imaginer.
   • Étape 1 (Étiquetage Pseudo) : Utilisation d'un modèle du monde et d'une politique initiale pour générer des étiquettes de "meilleur horizon" pour les trajectoires d'experts.
   • Étape 2 (Warm-up) : Entraînement supervisé conjoint de la politique d'action et d'un prédicteur d'horizon ($K$-head).
   • Étape 3 (Optimisation en ligne) : Utilisation d'un algorithme A2C (Advantage Actor-Critic) pour optimiser conjointement la politique d'action et la sélection de l'horizon, avec une fonction de récompense qui pénalise le coût computationnel excessif.

3. Contributions Clés

1. Conceptualisation du POIMDP : Introduction d'un nouveau cadre théorique intégrant explicitement les trajectoires futures imaginées dans le processus de décision des agents, au-delà du POMDP classique.
2. Cadre ITP et Mécanisme Adaptatif : Proposition d'une méthode unifiée qui combine les modèles du monde avec une sélection d'horizon dynamique, permettant une planification délibérée et efficace.
3. Validation Empirique : Démonstration que les deux variantes (ITP-I et ITP-R) surpassent significativement les méthodes de base (CoT, ReAct, RAP, SFT, etc.) sur des tâches complexes à long horizon, prouvant que l'adaptation de l'horizon est cruciale pour la performance.

4. Résultats Expérimentaux

Les expériences ont été menées sur quatre benchmarks représentatifs : ALFWorld (tâches domestiques), ScienceWorld (expériences scientifiques), WebShop (navigation web) et StableToolBench (utilisation d'outils).

• Performance Supérieure :
  • ITP-I (sans entraînement) améliore considérablement les performances zero-shot par rapport aux méthodes de prompting avancées (ReAct, RAP).
  • ITP-R (avec entraînement) surpasse toutes les méthodes basées sur l'entraînement (SFT, WKM, IWM). Par exemple, sur ALFWorld avec Qwen3-8B, ITP-R atteint un taux de réussite de 88,57 %, contre 82,86 % pour le meilleur concurrent (IWM).
• Efficacité Coût-Performance :
  • Les ablations montrent que l'approche adaptative surpasse largement les stratégies à horizon fixe (qui sont soit trop courtes, soit trop coûteuses) et les stratégies aléatoires.
  • ITP atteint un meilleur taux de réussite avec un budget de tokens normalisé (coût computationnel) inférieur.
• Robustesse aux Modèles du Monde :
  • Bien que la qualité du modèle du monde influence les résultats (surtout pour ITP-I), ITP-R est capable de réduire l'écart de performance entre différents modèles de base (Qwen, Llama, DeepSeek), démontrant une certaine agnosticité vis-à-vis de l'architecture du modèle du monde.
• Analyse de l'Horizon :
  • L'étude révèle que le taux de réussite atteint un pic à un horizon de prévision spécifique (ex: $k=3$ pour ALFWorld), au-delà duquel les erreurs d'accumulation du modèle du monde dégradent la performance. Le mécanisme adaptatif apprend naturellement à s'arrêter à cet horizon optimal.

5. Signification et Impact

Ce travail marque une avancée significative vers des agents autonomes plus délibératifs. En passant d'une réaction immédiate à une planification basée sur l'imagination future, ITP comble le fossé de l'ancrage superficiel.

• Théorique : Il propose un nouveau paradigme (POIMDP) pour intégrer la simulation mentale dans l'apprentissage par renforcement et le raisonnement des LLM.
• Pratique : Il offre une solution efficace pour les tâches complexes à long horizon où les erreurs sont coûteuses. La capacité à adapter dynamiquement l'effort de calcul (en imaginant plus ou moins loin selon le besoin) rend l'approche viable pour des applications réelles où l'efficacité est cruciale.
• Futur : Bien que les résultats soient prometteurs sur des environnements textuels, les auteurs soulignent que l'extension à des environnements multimodaux et robotiques réels nécessitera des optimisations supplémentaires pour gérer le bruit et les coûts d'inférence en temps réel.

En résumé, Imagine-then-Plan démontre que la capacité d'un agent à "imaginer" l'avenir de manière adaptative est un levier puissant pour améliorer la robustesse, la réussite et l'efficacité de l'apprentissage des agents autonomes.