Résumé technique : IMWM (Intuition Models Complement World Models for Latent Planning)

1. Énoncé du problème

Le papier traite d'une limitation fondamentale de l'apprentissage par renforcement (RL) fondé sur des modèles pour le contrôle à partir de pixels bruts : la dépendance à un modèle de monde (WM) latent appris associé à un planificateur à budget fini basé sur l'échantillonnage (par exemple, la méthode de l'entropie croisée, CEM). Bien que l'amélioration de la précision du modèle de monde soit l'approche standard pour accroître la performance de la planification, les auteurs soutiennent que cette intuition est insuffisante.

Le problème central identifié est que la recherche, et non la prédiction, peut être la contrainte limitante. Même si un planificateur possède un modèle de monde « parfait » (opérationnalisé comme un oracle utilisant la dynamique réelle de l'environnement pour les rollouts), il peut tout de même échouer à atteindre un objectif si le budget de recherche fini l'empêche de proposer une séquence d'actions réussie au sein du vaste espace d'actions. Le papier postule que le goulot d'étranglement réside dans le volume de proposition — la capacité du planificateur à tomber sur des candidats atteignant l'objectif — plutôt que dans la précision du prédicteur de l'état futur.

2. Diagnostic et base théorique

Avant de proposer une solution, les auteurs fournissent des preuves expérimentales et théoriques que la recherche est le point de défaillance :

• Diagnostic expérimental : Ils ont remplacé le prédicteur de l'état futur appris dans un planificateur latent standard par un déroulement idéalisé de la dynamique réelle de l'environnement (Oracle Dynamics). Malgré des dynamiques parfaites, le planificateur échouait encore de manière significative sur des tâches comme OGBench-Cube et Two-Room. L'analyse des épisodes échoués a révélé que dans la grande majorité des cas (par exemple, 98,9 % sur Two-Room), la population du CEM contenait zéro candidat atteignant l'objectif à n'importe quelle étape de replanification.
• Borne théorique : Les auteurs prouvent (Théorème A.1) que sous un budget fini de requêtes boîte noire contre toute surface de coût latente, la probabilité de succès est limitée par le volume de l'espace d'action couvert par les requêtes. Cette borne est indépendante de la manière dont le prédicteur a été entraîné, confirmant qu'un meilleur prédicteur ne peut surmonter un goulot d'étranglement de recherche si la distribution de proposition ne couvre pas la variété (manifold) du succès.

3. Méthodologie : IMWM

Pour remédier au goulot d'étranglement de la recherche sans nécessiter un meilleur modèle de monde, les auteurs proposent l'IMWM (Intuition Model + World Model). Ce cadre associe un modèle de monde latent gelé à un Modèle d'Intuition distinct, entraîné sur des démonstrations. Le modèle d'intuition agit comme un guide pour reconnaître les actions prometteuses, complétant plutôt que remplaçant le modèle de monde.

3.1 Le Modèle d'Intuition

Le modèle d'intuition ($D_{\psi}^{inv}$) est un score scalaire contrastif entraîné sur des fenêtres de démonstration à l'aide d'un objectif de type InfoNCE. Il prend un tuple de (état latent de départ, état latent cible, bloc d'actions) et produit un score de compatibilité.

• Conditionnement inverse : Contrairement au modèle de monde (qui prédit $s_{t+1}$ étant donné $s_t, a_t$), le modèle d'intuition évalue les actions étant donné les preuves de départ et de cible.
• Espace latent distinct : Il opère dans un espace latent distinct ($z_I$) découplé de l'espace de prédiction directe du modèle de monde ($z_L$), évitant ainsi l'accumulation des erreurs de déroulement.
• Entraînement : Il est entraîné pour distinguer les blocs d'actions de démonstration des échantillons négatifs au sein du même contexte de tâche.

3.2 Trois composants de collaboration légers

Le modèle d'intuition interagit avec le planificateur via trois mécanismes spécifiques :

1. Initialisation par récupération (Retrieval Initialization) :

   • Avant que la fonction de coût n'évalue les candidats, le planificateur récupère un bloc d'actions à partir d'une banque fixe de démonstrations.
   • La récupération est basée sur la similitude cosinus entre les latents actuels (départ, cible) et les clés de démonstration stockées dans l'encodeur d'intuition dans l'espace latent.
   • La distribution de proposition du CEM est initialisée comme une Gaussienne isotrope centrée sur ce bloc d'actions récupéré, biaisant la recherche vers des régions soutenues par les démonstrations.

2. Coût Hybride :

   • Le planificateur minimise une fonction de coût composite $J$ qui combine le score d'intuition et l'erreur de déroulement du modèle de monde.
   • $J = \alpha_{inv} \cdot zS(-D_{\psi}^{inv}) + \beta_{eff} \cdot zS(\text{MSE}(\text{rollout}, \text{goal}))$.
   • Les deux termes sont normalisés (z-scored) à travers l'ensemble des candidats pour gérer les échelles incompatibles. Les poids $\alpha$ et $\beta$ déterminent la dépendance envers l'intuition ou le modèle de monde.

3. Porte de Fiabilité (Reliability Gate) :

   • Puisque l'intuition n'est pas universellement utile (par exemple, si la banque de démonstrations manque d'exemples pertinents), une porte décide par tâche/cellule du degré de confiance accordé au modèle d'intuition.
   • La porte calcule deux diagnostics à partir du bloc récupéré : une marge de score ($r_{inv}$) par rapport aux blocs neutres et une autocorrélation de décalage 1 ($r_{lag}$).
   • Basée sur des seuils gelés, la porte sélectionne l'une des trois "recettes" préspécifiées :
     • Hybride de Confiance (Trusted Hybrid) : Forte dépendance à la fois à l'intuition et au modèle de monde.
     • Dominante Intuition (Intuition-Dominant) : Forte dépendance à l'intuition, faible dépendance au modèle de monde.
     • Repli Uniquement Forward (Forward-Only Fallback) : L'intuition est désactivée ; le planificateur revient à la configuration standard du seul modèle de monde.
   • Crucialement, cette sélection est automatique basée sur les diagnostics, et non ajustée manuellement par tâche.

4. Contributions Clés

1. Diagnostic du planning latent à budget CEM fini : Le papier démontre expérimentalement et théoriquement que même avec un modèle de monde aux dynamiques oracle, les planificateurs à budget fini échouent en raison de limitations de recherche (volume de proposition) plutôt que d'erreurs de prédiction.
2. La méthode IMWM : Une architecture novatrice qui intègre un modèle d'intuition contrastif à un modèle de monde via l'initialisation par récupération, le coût hybride et la porte de fiabilité.
3. Gains empiriques : Améliorations significatives par rapport aux bases de référence uniquement basées sur le modèle de monde sur quatre tâches de franchissement de but basées sur des pixels, sans réglage par tâche de l'algorithme central.

5. Résultats Expérimentaux

La méthode a été évaluée sur une grille de 48 cellules (4 tâches $\times$ 12 graines) utilisant la plateforme stable-worldmodel. Les tâches étaient : Two-Room, Reacher, Push-T et OGBench-Cube.

• Two-Room : IMWM a atteint un taux de succès de 99,2 %, soit une amélioration de +11,5 points de pourcentage (pp) par rapport à la base (87,7 %).
• OGBench-Cube : IMWM a atteint un taux de succès de 94,7 %, une amélioration massive de +28,5 pp par rapport à la base (66,2 %).
• Push-T : IMWM a atteint 92,7 %, soit une amélioration de +2,8 pp (base 89,8 %).
• Reacher : IMWM a atteint 83,8 %, un résultat quasi identique à la base (83,2 %, +0,7 pp). Dans ce cas, la Porte de Fiabilité a correctement dirigé vers la recette de "repli uniquement forward", car le signal d'intuition n'était pas bénéfique.

Études d'ablation :

• Initialisation par récupération : Activer uniquement la récupération a considérablement boosté la performance sur les tâches limitées par la recherche (par exemple, OGBench-Cube passant de 66,2 % à 88,3 %) mais a pénalisé les performances sur les tâches où le modèle de monde était déjà fort (Reacher), validant la nécessité de la porte.
• Coût Hybride : La configuration complète (les deux coûts d'intuition et de modèle de monde) a surpassé l'utilisation de l'un ou l'autre terme seul.
• Porte de Fiabilité : Aucune recette fixe unique n'était optimale pour toutes les tâches. La porte adaptative a sélectionné avec succès la meilleure recette pour chaque tâche automatiquement, égalant la performance d'un scénario optimisé fait à la main.

6. Signification et Revendications

Le papier affirme que le principal goulot d'étranglement de la planification par modèle de monde latent est souvent la stratégie de recherche, et non la précision prédictive du modèle. En introduisant un mécanisme d'« intuition » dérivé des démonstrations, l'IMWM concentre efficacement le budget de recherche sur les régions prometteuses de l'espace d'actions, contournant ainsi les limitations de volume du l'échantillonnage aveugle.

Les auteurs soulignent que :

• Les améliorations ne sont pas dues à un meilleur modèle de monde ou à un réglage des hyperparamètres par tâche ; la famille algorithmique et les artefacts entraînés sont fixes ; seule la sélection de la recette par la porte varie selon les diagnostics pré-calculés.
• L'approche est complémentaire : le modèle d'intuition ne remplace pas le modèle de monde mais le guide.
• Limitations : La porte de fiabilité est conçue à la main avec des seuils fixes (bien qu'interprétables), les résultats sont limités à une famille de benchmarks de tâches basées sur des pixels, et la méthode nécessite actuellement des démonstrations réussies pour construire la banque d'intuition. Les auteurs ne revendiquent pas de transfert au contrôle de robots réels ou à des benchmarks plus larges sans travail supplémentaire.

En conclusion, l'IMWM démontre que coupler un modèle de monde avec un modèle d'intuition basé sur la récupération peut résoudre le mode d'échec de "recherche" de la planification latente, atteignant des performances de pointe sur des tâches spécifiques de franchissement de but là où la planification par pur modèle de monde peine.