Probabilistic Dreaming for World Models

Cette étude propose des améliorations probabilistes au modèle Dreamer permettant l'exploration parallèle d'états latents et le maintien d'hypothèses distinctes pour des futurs mutuellement exclusifs, ce qui se traduit par de meilleures performances et une variance réduite sur le domaine MPE SimpleTag.

L'essentiel

🌌 Le Rêve de l'IA : Apprendre en dormant

Imaginez que vous voulez apprendre à jouer à un jeu vidéo très difficile. Vous avez deux options :

1. La méthode classique : Vous jouez, vous vous faites tuer, vous recommencez. C'est long, fatiguant et vous apprenez lentement.
2. La méthode "Dreamer" (Rêveur) : Au lieu de jouer, vous fermez les yeux et vous rêvez de parties futures. Votre cerveau simule ce qui pourrait arriver, vous testez des stratégies dans votre tête, et vous apprenez sans bouger un muscle. C'est ce qu'on appelle l'apprentissage par "imagination latente".

Le papier que nous allons explorer parle d'une nouvelle version de ce rêveur, appelée ProbDreamer. L'auteur, Gavin Wong, se demande : "Et si au lieu de rêver d'un seul scénario, on rêvait de plusieurs en même temps ?"

---

🧠 Le Problème du Rêveur Classique

Le modèle précédent (Dreamer) est intelligent, mais il a un défaut de conception : il est trop sûr de lui.

Imaginez que vous êtes un chat poursuivi par un chien. Le chien peut soit vous chasser (courir derrière vous), soit vous intercepter (couper votre chemin).

• Le rêveur classique essaie de prédire l'avenir en faisant une moyenne.
• Il se dit : "Le chien va faire 50% de chasse et 50% d'interception."
• Résultat de son rêve ? Il imagine un chien qui fait un mouvement bizarre, mi-chasse, mi-interception, qui n'existe pas dans la réalité. C'est comme si le chien essayait de faire les deux en même temps et se figeait.
• Dans le jeu, cela signifie que l'IA se fige un instant, panique, et se fait attraper.

✨ La Solution : Le "Rêveur à Particules"

Pour régler ce problème, Gavin Wong propose d'utiliser une technique appelée Filtre à Particules. Voici l'analogie :

Au lieu d'avoir un seul rêveur qui fait une moyenne, imaginez que vous avez un petit groupe d'amis (des particules) qui rêvent tous en même temps pour vous.

• Particule A rêve : "Le chien va me chasser ! Je vais courir à gauche."
• Particule B rêve : "Le chien va m'intercepter ! Je vais courir à droite."

Chaque ami garde son hypothèse bien distincte. Ils ne mélangent pas leurs idées. Ils gardent deux scénarios clairs : "Gauche" et "Droite".

• Avantage : Quand le chien choisit réellement sa stratégie, l'IA a déjà préparé deux réponses possibles. Elle ne se fige pas. Elle réagit vite.

🚦 Comment ça marche en pratique ?

Le papier décrit trois innovations principales pour rendre ce rêve collectif efficace :

1. Le Rêve en Parallèle : Au lieu de simuler une seule ligne du temps, le modèle simule plusieurs lignes de temps (K particules) en même temps. C'est comme avoir plusieurs versions de vous-même qui testent différentes stratégies simultanément.
2. La Recherche de Chemins (Beam Search) : Pour chaque ami (particule), on imagine plusieurs actions possibles. Si la particule A pense "courir à gauche", on imagine aussi "courir à gauche vite" et "courir à gauche lentement". Cela crée une forêt d'arbres de possibilités.
3. Le Tri par "Énergie Libre" : Comme on ne peut pas vérifier si les rêves sont vrais (on est en train de dormir !), il faut un système pour garder les meilleurs rêves et jeter les mauvais. Le modèle utilise une formule mathématique (l'énergie libre) qui dit : "Garde les rêves qui rapportent beaucoup de points ET ceux qui sont les plus incertains (les plus curieux)."

🏆 Les Résultats : Qui gagne ?

Les chercheurs ont testé cela sur un jeu simple où un agent doit échapper à trois prédateurs qui changent de stratégie au hasard.

• Le Vainqueur : La version "Lite" de ProbDreamer (avec seulement 2 amis qui rêvent) a gagné.
• La Performance : Elle a obtenu un score 4,5 % meilleur que le rêveur classique.
• La Robustesse : C'est le point le plus important : ses performances sont beaucoup plus stables (28 % de moins de variations). Elle ne fait pas de "crises de panique" comme l'ancienne version.

L'analogie finale :
Le rêveur classique est comme un conducteur qui, face à un virage, calcule la moyenne entre "tourner à gauche" et "tourner à droite", et finit par foncer dans le mur.
Le nouveau rêveur (ProbDreamer) est comme un conducteur qui a deux copilotes : l'un crie "Tourne à gauche !", l'autre crie "Tourne à droite !". Le conducteur écoute les deux, et dès que la route se précise, il sait exactement quoi faire.

⚠️ Les Limites et le Futur

Ce n'est pas encore parfait. Les chercheurs ont remarqué deux problèmes :

1. Le nombre d'amis (particules) : Avec trop d'amis (trop de particules), le système devient confus et commence à rêver n'importe quoi (du bruit). Il faut trouver le juste milieu.
2. Le tri des rêves : Le système pour choisir les "meilleurs" rêves (basé sur l'incertitude) ne fonctionnait pas très bien car les "copilotes" finissaient tous par se mettre d'accord trop vite, perdant leur diversité.

Conclusion :
Ce papier montre que pour qu'une IA apprenne vraiment bien, elle ne doit pas juste "moyenner" le futur. Elle doit maintenir plusieurs hypothèses contradictoires en même temps, comme le fait un humain qui imagine plusieurs issues possibles avant d'agir. C'est un pas de géant vers des intelligences artificielles plus curieuses, plus robustes et capables de mieux naviguer dans un monde imprévisible.

Résumé technique

1. Problématique

L'apprentissage par renforcement (RL) basé sur des modèles, et plus particulièrement l'approche « Dreamer », permet aux agents d'apprendre en imaginant des trajectoires futures via un modèle du monde, réduisant ainsi le besoin d'interactions coûteuses avec l'environnement réel. Cependant, l'article identifie deux limites majeures dans les méthodes actuelles (Dreamer v3/v4) :

1. Limitation de l'exploration : Bien que le modèle apprenne une distribution complète d'états latents, il ne tire qu'un seul état pour dérouler une trajectoire imaginaire unique. Cela limite l'exploration de la diversité des causes possibles durant l'entraînement.
2. Le problème de la multimodalité : Les latents continus (Gaussiens) sont préférés pour leurs propriétés de gradient lisses, mais une Gaussienne unimodale a tendance à « moyenner » des alternatives mutuellement exclusives (par exemple, fusionner un chemin « gauche » et un chemin « droite » en un chemin « milieu » impossible). À l'inverse, les latents discrets récents gèrent la multimodalité mais perdent ces propriétés de gradient.

L'objectif est donc d'améliorer le processus d'imagination latente en intégrant des méthodes probabilistes pour maintenir des hypothèses distinctes tout en conservant les avantages des latents continus.

2. Méthodologie

Les auteurs proposent une architecture nommée ProbDreamer, basée sur Dreamer-v3 mais modifiée avec trois innovations clés :

• Filtre à Particules (Particle Filter) : Au lieu de tirer un seul état latent, l'agent maintient un ensemble de $K$ particules $\{h^k_t, z^k_t\}$ qui suivent la distribution a priori. Cela permet d'avoir $K « rêves » parallèles$ à chaque étape d'entraînement. Chaque transition reste gaussienne, mais la croyance globale sur les états latents devient une distribution empirique sur les particules, capable d'approximer des croyances multimodales complexes.
• Recherche en Faisceau Latent (Latent Beam Search) : Pour propager chaque particule, l'agent bifurque explicitement chaque particule en $N$ actions candidates échantillonnées depuis la politique. Cela crée $K \times N$ branches qui sont propagées via le modèle du monde.
• Minimisation de l'Énergie Libre (Free Energy) : Pour élaguer (pruner) les trajectoires imaginées sans observations réelles, les auteurs utilisent un critère combinant la récompense prédite et l'incertitude épistémique (mesurée par la variance d'un ensemble de modèles a priori). La fonction de score est :
  $$F^k_t = V_\phi(h^k_t, z^k_t) + \beta \cdot \sigma^2_{ens}$$
  où $V_\phi$ est la valeur prédite par le critique et $\sigma^2_{ens}$ la variance de l'ensemble. L'objectif est de maximiser la récompense tout en explorant les situations à forte incertitude.

L'évaluation est réalisée sur le domaine MPE SimpleTag, un jeu prédateur-proie où les prédateurs changent stochastiquement de stratégie (« Chasse » vs « Intercept »), créant un espace d'états bimodal.

3. Contributions Clés

1. Résolution du problème de moyennage : L'utilisation d'un filtre à particules permet de maintenir des hypothèses concurrentes distinctes (ex: stratégie de chasse vs interception) sans forcer le modèle à converger vers une moyenne impossible, tout en conservant les latents continus.
2. Exploration parallèle : L'approche permet l'exploration simultanée de multiples états latents et branches d'actions, augmentant la couverture de l'espace des causes.
3. Preuve de concept : L'article démontre que des méthodes probabilistes non paramétriques peuvent améliorer l'efficacité et la robustesse du RL basé sur des modèles, tout en identifiant les limites actuelles des mécanismes d'élagage.

4. Résultats

Les expériences ont été menées avec une recherche d'hyperparamètres rigoureuse (Optimisation Bayésienne) sur 6 configurations finales.

• Performance et Robustesse : La variante « Lite » ProbDreamer (2 particules, sans recherche en faisceau, $K=2, N=1$) surpasse systématiquement le modèle de base (BaseDreamer) sur 4 des 5 graines aléatoires.
  • Amélioration du score : +4,5 % par rapport au modèle de base.
  • Réduction de la variance : 28 % de variance en moins dans les retours d'épisode, indiquant une politique plus robuste.
• Analyse Comportementale : L'agent ProbDreamer réagit rapidement aux changements de stratégie des prédateurs, tandis que le modèle de base (Gaussien) a tendance à « geler » momentanément, signe d'un biais de moyenne paralysant.
• Échec des configurations complexes : Les modèles « Full » (avec recherche en faisceau et $K$ élevé, ex: $K=8$) ont montré une dégradation sévère des performances.
  • Saturation des particules : Un $K$ trop élevé (au-delà du nombre de stratégies réelles) introduit du bruit.
  • Élagage inefficace : L'élagage basé uniquement sur une fonction de valeur apprise (sans observations réelles) conduit à sélectionner des trajectoires irréalistes (hallucinations optimistes) lorsque le critique est bruyant.
  • Effondrement de l'ensemble : Les membres de l'ensemble utilisé pour estimer l'incertitude épistémique ont rapidement convergé vers des prédictions identiques, rendant le terme de curiosité inefficace.

5. Signification et Perspectives

Ce travail valide l'hypothèse que l'échantillonnage probabiliste multiple d'états latents (« particules ») est supérieur à l'échantillonnage unique pour gérer l'ambiguïté multimodale dans les modèles du monde.

Limitations identifiées et travaux futurs :

• Complexité de l'environnement : Le domaine SimpleTag étant partiellement observable mais relativement simple, un filtre à 2 particules suffisait. Il faudra tester comment le nombre optimal de particules ($K$) évolue avec la complexité environnementale (environnements chaotiques ou partiellement observables).
• Gestion de l'incertitude : Le principal goulot d'étranglement reste l'absence d'observations réelles pour corriger les « rêves ». Les auteurs suggèrent de développer des architectures capables de capturer intrinsèquement l'incertitude épistémique (par exemple via des approximations bayésiennes plus riches comme le Dropout Monte-Carlo ou des ensembles diversifiés) pour permettre un équilibre autonome entre exploration et exploitation, mimant ainsi l'apprentissage curieux humain.

En conclusion, bien que l'élagage actif des trajectoires imaginées pose encore des défis, l'approche par filtre à particules offre une voie prometteuse pour les modèles du monde non paramétriques.