C’était en 2018, lorsque l’idée d’apprentissage par renforcement dans le contexte d’un modèle de monde neuronal a été introduite pour la première fois, et bientôt, ce principe fondamental a été appliqué aux modèles de monde. Certains des modèles les plus importants qui mettent en œuvre l’apprentissage par renforcement étaient le cadre Dreamer, qui a introduit l’apprentissage par renforcement à partir de l’espace latent d’un modèle d’espace d’état récurrent. Le DreamerV2 a démontré que l’utilisation de latents discrets pourrait entraîner une réduction des erreurs de composition, et le cadre DreamerV3 a pu atteindre des performances humaines sur une série de tâches dans différents domaines avec des hyperparamètres fixes.

De plus, des parallèles peuvent être établis entre les modèles de génération d’images et les modèles de monde, indiquant que les progrès réalisés dans les modèles de vision génératifs pourraient être reproduits pour bénéficier aux modèles de monde. Depuis que l’utilisation de transformateurs dans les cadres de traitement du langage naturel a gagné en popularité, les cadres DALL-E et VQGAN ont émergé. Les cadres ont mis en œuvre des auto-encodeurs discrets pour convertir les images en jetons discrets et ont pu construire des modèles génératifs de texte à image très puissants et efficaces en exploitant les capacités de modélisation de séquence des transformateurs autorégressifs. En même temps, les modèles de diffusion ont gagné en popularité, et aujourd’hui, les modèles de diffusion se sont établis comme un paradigme dominant pour la génération d’images à haute résolution. Grâce aux capacités offertes par les modèles de diffusion et l’apprentissage par renforcement, des tentatives sont faites pour combiner les deux approches, dans le but de profiter de la flexibilité des modèles de diffusion en tant que modèles de trajectoire, modèles de récompense, planificateurs et politiques pour l’augmentation de données dans l’apprentissage par renforcement hors ligne.

Les modèles de monde offrent une méthode prometteuse pour former des agents d’apprentissage par renforcement de manière sûre et efficace. Traditionnellement, ces modèles utilisent des séquences de variables latentes discretes pour simuler la dynamique de l’environnement. Cependant, cette compression peut négliger les détails visuels cruciaux pour l’apprentissage par renforcement. En même temps, les modèles de diffusion ont gagné en popularité pour la génération d’images, remettant en question les méthodes traditionnelles qui utilisent des latents discrets. Inspirés par ce changement, dans cet article, nous allons parler de DIAMOND (DIffusion As a Model Of eNvironment Dreams), un agent d’apprentissage par renforcement formé dans un modèle de monde de diffusion. Nous allons explorer les choix de conception nécessaires pour rendre la diffusion adaptée à la modélisation du monde et montrer que des détails visuels améliorés conduisent à de meilleures performances d’agent. DIAMOND établit un nouveau benchmark sur le test compétitif Atari 100k, atteignant un score humain normalisé moyen de 1,46, le plus élevé pour les agents formés entièrement dans un modèle de monde.

DIAMOND : La diffusion comme un modèle de rêves d’environnement

Les modèles de monde ou les modèles génératifs d’environnement émergent comme l’un des composants les plus importants pour les agents génératifs pour planifier et raisonner sur leurs environnements. Bien que l’utilisation de l’apprentissage par renforcement ait obtenu un succès considérable ces dernières années, les modèles qui mettent en œuvre l’apprentissage par renforcement sont connus pour être inefficaces en termes d’échantillons, ce qui limite considérablement leurs applications dans le monde réel. D’un autre côté, les modèles de monde ont démontré leur capacité à former des agents d’apprentissage par renforcement de manière efficace dans divers environnements avec une efficacité d’échantillonnage considérablement améliorée, permettant au modèle d’apprendre à partir d’expériences du monde réel. Les cadres de modélisation du monde récents modélisent généralement la dynamique de l’environnement comme une séquence de variables latentes discretes, avec le modèle qui discrétise l’espace latent pour éviter les erreurs de composition sur des horizons de temps multiples. Bien que cette approche puisse donner des résultats substantiels, elle est également associée à une perte d’information, entraînant une perte de qualité de reconstruction et de généralité. La perte d’information peut devenir un obstacle important pour les scénarios du monde réel qui nécessitent que l’information soit bien définie, comme la formation de véhicules autonomes. Dans de telles tâches, de petits changements ou des détails dans les entrées visuelles, tels que la couleur du feu de circulation ou l’indicateur de direction du véhicule avant, peuvent changer la politique d’un agent. Bien que l’augmentation du nombre de latents discrets puisse aider à éviter la perte d’information, elle augmente considérablement les coûts de calcul.

De plus, ces dernières années, les modèles de diffusion ont émergé comme l’approche dominante pour les cadres de génération d’images de haute qualité, car les cadres construits sur les modèles de diffusion apprennent à inverser un processus de bruit, et concurrencent directement certaines des approches les plus établies qui modélisent des jetons discrets, offrant ainsi une alternative prometteuse pour éliminer le besoin de discrétisation dans la modélisation du monde. Les modèles de diffusion sont connus pour leur capacité à être facilement conditionnés et à modéliser de manière flexible des distributions complexes et multimodales sans effondrement de mode. Ces attributs sont cruciaux pour la modélisation du monde, car la conditionnalisation permet à un modèle de monde de refléter avec précision les actions d’un agent, conduisant à une attribution de crédit plus fiable. De plus, la modélisation de distributions multimodales offre une plus grande diversité de scénarios de formation pour l’agent, améliorant ainsi ses performances globales.

En nous appuyant sur ces caractéristiques, DIAMOND, (DIffusion As a Model Of eNvironment Dreams), un agent d’apprentissage par renforcement formé dans un modèle de monde de diffusion. Le cadre DIAMOND fait des choix de conception soigneux pour s’assurer que son modèle de monde de diffusion reste efficace et stable sur de longs horizons de temps. Le cadre fournit une analyse qualitative pour démontrer l’importance de ces choix de conception. DIAMOND établit un nouveau record avec un score humain normalisé moyen de 1,46 sur le benchmark établi d’Atari 100k, le plus élevé pour les agents formés entièrement dans un modèle de monde. Le fonctionnement dans l’espace d’images permet au modèle de monde de diffusion de DIAMOND de substituer sans effort l’environnement, offrant ainsi une meilleure compréhension du comportement du modèle de monde et de l’agent. Notamment, les performances améliorées dans certains jeux sont attribuées à une meilleure modélisation des détails visuels critiques. Le cadre DIAMOND modélise l’environnement comme un processus de décision markovien partiellement observable standard avec un ensemble d’états, un ensemble d’actions discretes et un ensemble d’observations d’images. Les fonctions de transition décrivent la dynamique de l’environnement, et la fonction de récompense mappe les transitions à des récompenses scalaires. La fonction d’observation décrit les probabilités d’observation et émet des observations d’images, qui sont ensuite utilisées par les agents pour voir les environnements, puisqu’ils ne peuvent pas accéder directement aux états. L’objectif principal de l’approche était d’obtenir une politique qui mappe les observations aux actions avec la tentative de maximiser la récompense attendue avec un facteur de décompte. Les modèles de monde sont des modèles génératifs de l’environnement, et les modèles de monde peuvent être utilisés pour créer des environnements simulés pour former des agents d’apprentissage par renforcement dans l’environnement réel et former des agents d’apprentissage par renforcement dans l’environnement du modèle de monde. La figure 1 démontre l’imagination déroulée du cadre DIAMOND au fil du temps.

DIAMOND : Méthodologie et architecture

Au cœur, les modèles de diffusion sont une classe de modèles génératifs qui génèrent un échantillon en inversant le processus de bruit, et tirent fortement leur inspiration de la thermodynamique non équilibre. Le cadre DIAMOND considère un processus de diffusion indexé par une variable de temps continue avec des marginales et des conditions limites correspondantes avec une distribution a priori non structurée tractable. De plus, pour obtenir un modèle génératif, qui mappe du bruit aux données, le cadre DIAMOND doit inverser le processus, avec le processus d’inversion étant également un processus de diffusion, s’exécutant à rebours dans le temps. De plus, à tout moment donné, il n’est pas trivial d’estimer la fonction de score, car le cadre DIAMOND n’a pas accès à la fonction de score réelle, et le modèle surmonte cet obstacle en mettant en œuvre l’objectif de correspondance de score, une approche qui permet à un cadre de former un modèle de score sans connaître la fonction de score sous-jacente. Le modèle de diffusion basé sur le score fournit un modèle génératif inconditionnel. Cependant, un modèle génératif conditionnel de la dynamique de l’environnement est requis pour servir de modèle de monde, et pour servir cet objectif, le cadre DIAMOND examine le cas général de l’approche POMDP, dans laquelle le cadre peut utiliser les observations et les actions passées pour approximer l’état markovien inconnu. Comme le démontre la figure 1, le cadre DIAMOND utilise cet historique pour conditionner un modèle de diffusion, pour estimer et générer la prochaine observation directement. Bien que le cadre DIAMOND puisse en théorie recourir à tout solveur SDE ou ODE, il existe un compromis entre le nombre d’évaluations de fonction et la qualité de l’échantillon qui affecte considérablement le coût d’inférence des modèles de diffusion.

En nous appuyant sur ces apprentissages, examinons maintenant la réalisation pratique du cadre DIAMOND d’un modèle de monde basé sur la diffusion, y compris les coefficients de dérive et de diffusion correspondant à un choix particulier d’approche de diffusion. Au lieu d’opter pour le DDPM, un candidat naturellement adapté à la tâche, le cadre DIAMOND se base sur la formulation EDM, et considère un noyau de perturbation avec une fonction à valeur réelle de temps de diffusion appelée le calendrier de bruit. Le cadre sélectionne les préconditionneurs pour conserver la variance d’entrée et de sortie pour tout niveau de voix. La formation du réseau combine le signal et le bruit de manière adaptative en fonction du niveau de dégradation, et lorsque le bruit est faible, et que la cible devient la différence entre le signal propre et le signal perturbé, c’est-à-dire le bruit gaussien ajouté. Intuitivement, cela empêche l’objectif de formation de devenir trivial dans le régime de faible bruit. En pratique, cet objectif est à haute variance aux extrêmes du calendrier de bruit, donc le modèle échantillonne le niveau de bruit à partir d’une distribution log-normale choisie empiriquement afin de concaténer la formation autour des régions de bruit moyen. Le cadre DIAMOND utilise un composant U-Net 2D standard pour le champ de vecteur, et conserve un tampon d’observations et d’actions passées que le cadre utilise pour se conditionner. Le cadre DIAMOND concatène ensuite ces observations passées à la prochaine observation bruyante, et les actions d’entrée via des couches de normalisation de groupe adaptative dans les blocs résiduels de l’U-Net.

DIAMOND : Expériences et résultats

Pour une évaluation complète, le cadre DIAMOND opte pour le benchmark Atari 100k. Le benchmark Atari 100k se compose de 26 jeux conçus pour tester une large gamme de capacités d’agent. Dans chaque jeu, un agent est limité à 100 000 actions dans l’environnement, ce qui équivaut approximativement à 2 heures de jeu humain, pour apprendre le jeu avant l’évaluation. Pour comparaison, les agents Atari non contraints s’entraînent généralement pendant 50 millions d’étapes, représentant une augmentation de 500 fois de l’expérience. Nous avons formé DIAMOND à partir de zéro en utilisant 5 graines aléatoires pour chaque jeu. Chaque exécution d’entraînement nécessitait environ 12 Go de VRAM et prenait environ 2,9 jours sur une seule carte Nvidia RTX 4090, ce qui équivaut à 1,03 année de GPU au total. Le tableau suivant fournit le score pour tous les jeux, la moyenne et la moyenne interquartile des scores humains normalisés.

En suivant les limites des estimations ponctuelles, le cadre DIAMOND fournit une confiance bootstrap stratifiée dans la moyenne et la moyenne interquartile des scores humains normalisés, ainsi que des profils de performance et des métriques supplémentaires, comme résumé dans la figure suivante.

Les résultats montrent que DIAMOND se comporte de manière exceptionnelle sur l’ensemble du benchmark, surpassant les joueurs humains dans 11 jeux et atteignant un score humain normalisé moyen de 1,46, établissant un nouveau record pour les agents formés entièrement dans un modèle de monde. De plus, la moyenne interquartile de DIAMOND est comparable à STORM et dépasse tous les autres références. DIAMOND excelle dans les environnements où la capture de petits détails est cruciale, tels que Asterix, Breakout et RoadRunner. De plus, comme mentionné précédemment, le cadre DIAMOND a la flexibilité de mettre en œuvre n’importe quel modèle de diffusion dans son pipeline, bien qu’il opte pour l’approche EDM, il aurait été naturel de choisir le modèle DDPM, car il est déjà mis en œuvre dans de nombreuses applications de génération d’images. Pour comparer l’approche EDM à l’implémentation DDPM, le cadre DIAMOND forme les deux variantes avec la même architecture de réseau sur le même jeu de données statique partagé avec plus de 100 000 cadres collectés avec une politique d’expert. Le nombre d’étapes de débruitage est directement lié au coût d’inférence du modèle de monde, et donc moins d’étapes réduiront le coût de formation d’un agent sur des trajectoires imaginées. Pour nous assurer que notre modèle de monde reste comparable en termes de coût de calcul avec d’autres références, telles que IRIS qui nécessite 16 évaluations de fonction par pas de temps, nous visons à utiliser au plus quelques étapes de débruitage, de préférence moins. Cependant, définir le nombre d’étapes de débruitage trop bas peut dégrader la qualité visuelle, entraînant des erreurs de composition. Pour évaluer la stabilité des différentes variantes de diffusion, nous affichons des trajectoires imaginées de manière autonome jusqu’à t = 1000 pas de temps dans la figure suivante, en utilisant différents nombres d’étapes de débruitage n ≤ 10.

Nous observons que l’utilisation de DDPM (a) dans ce régime entraîne des erreurs de composition graves, provoquant le modèle de monde à dériver rapidement hors de la distribution. En revanche, le modèle de monde de diffusion basé sur EDM (b) reste beaucoup plus stable sur de longs horizons de temps, même avec une seule étape de débruitage. Les trajectoires imaginées avec des modèles de monde de diffusion basés sur DDPM (gauche) et EDM (droite) sont montrées. L’observation initiale à t = 0 est la même pour les deux, et chaque ligne correspond à un nombre décroissant d’étapes de débruitage n. Nous observons que la génération basée sur DDPM souffre d’erreurs de composition, avec un nombre plus petit d’étapes de débruitage entraînant une accumulation d’erreurs plus rapide. En revanche, le modèle de monde de diffusion de DIAMOND reste beaucoup plus stable, même pour n = 1. La prédiction optimale en une seule étape est l’attente sur les reconstructions possibles pour une entrée bruyante donnée, qui peut être hors de la distribution si la distribution postérieure est multimodale. Bien que certaines parties du jeu, comme Breakout, aient des transitions déterministes qui peuvent être modélisées avec précision avec une seule étape de débruitage, d’autres jeux présentent une observabilité partielle, entraînant des distributions d’observation multimodales. Dans ces cas, un solveur itératif est nécessaire pour guider la procédure d’échantillonnage vers un mode spécifique, comme illustré dans le jeu Boxing dans la figure suivante. Par conséquent, le cadre DIAMOND définit n = 3 dans toutes nos expériences.

La figure ci-dessus compare l’échantillonnage en une seule étape (ligne du haut) et l’échantillonnage en plusieurs étapes (ligne du bas) dans Boxing. Les mouvements du joueur noir sont imprévisibles, provoquant l’échantillonnage en une seule étape pour interpoler entre les résultats possibles, entraînant des prédictions floues. En revanche, l’échantillonnage en plusieurs étapes produit une image claire en guidant la génération vers un mode spécifique. Intéressant, puisque la politique contrôle le joueur blanc, ses actions sont connues du modèle de monde, éliminant ainsi l’ambiguïté. Ainsi, les deux échantillonnages en une seule étape et en plusieurs étapes prédisent correctement la position du joueur blanc.

Dans la figure ci-dessus, les trajectoires imaginées par DIAMOND présentent généralement une qualité visuelle plus élevée et sont plus fidèles à l’environnement réel par rapport à celles imaginées par IRIS. Les trajectoires générées par IRIS contiennent des incohérences visuelles entre les cadres (soulignées par des boîtes blanches), telles que des ennemis affichés comme des récompenses et vice-versa. Bien que ces incohérences puissent n’affecter que quelques pixels, elles peuvent avoir un impact significatif sur l’apprentissage par renforcement. Par exemple, un agent vise généralement à cibler les récompenses et à éviter les ennemis, donc ces petites disparités visuelles peuvent rendre plus difficile l’apprentissage d’une politique optimale. La figure montre des cadres consécutifs imaginés avec IRIS (gauche) et DIAMOND (droite). Les boîtes blanches mettent en évidence les incohérences entre les cadres, qui n’apparaissent que dans les trajectoires générées avec IRIS. Dans Asterix (ligne du haut), un ennemi (orange) devient une récompense (rouge) dans le deuxième cadre, puis revient à un ennemi dans le troisième, et à nouveau à une récompense dans le quatrième. Dans Breakout (ligne du milieu), les briques et le score sont incohérents entre les cadres. Dans Road Runner (ligne du bas), les récompenses (petits points bleus sur la route) sont rendues de manière incohérente entre les cadres. Ces incohérences n’apparaissent pas avec DIAMOND. Dans Breakout, le score est mis à jour de manière fiable par +7 lorsque une brique rouge est cassée.

Conclusion

Dans cet article, nous avons parlé de DIAMOND, un agent d’apprentissage par renforcement formé dans un modèle de monde de diffusion. Le cadre DIAMOND fait des choix de conception soigneux pour s’assurer que son modèle de monde de diffusion reste efficace et stable sur de longs horizons de temps. Le cadre fournit une analyse qualitative pour démontrer l’importance de ces choix de conception. DIAMOND établit un nouveau record avec un score humain normalisé moyen de 1,46 sur le benchmark établi d’Atari 100k, le plus élevé pour les agents formés entièrement dans un modèle de monde. Le fonctionnement dans l’espace d’images permet au modèle de monde de diffusion de DIAMOND de substituer sans effort l’environnement, offrant ainsi une meilleure compréhension du comportement du modèle de monde et de l’agent. Notamment, les performances améliorées dans certains jeux sont attribuées à une meilleure modélisation des détails visuels critiques. Le cadre DIAMOND modélise l’environnement comme un processus de décision markovien partiellement observable standard avec un ensemble d’états, un ensemble d’actions discretes et un ensemble d’observations d’images. Les fonctions de transition décrivent la dynamique de l’environnement, et la fonction de récompense mappe les transitions à des récompenses scalaires.