Simulus: Combining Improvements in Sample-Efficient World Model Agents

Résumé Technique : Simulus – Combinaison d'Améliorations pour des Agents de Modèles du Monde Économes en Échantillons

Énoncé du Problème
Les modèles du monde (WM) représentent une frontière dans l'apprentissage par renforcement (RL) économe en échantillons, pourtant leur potentiel reste souvent inexploité en raison de la complexité d'intégration d'améliorations prometteuses. Bien que des avancées individuelles existent, leur combinaison requiert une expertise significative et des efforts d'ingénierie. De plus, les modèles du monde basés sur des tokens (TBWM) existants sont largement restreints aux observations d'images et aux actions discrètes, limitant leur applicabilité au contrôle continu et aux paramètres multi-modaux. En outre, des composants tels que la motivation intrinsèque (pour la réduction de l'incertitude épistémique) et la mémoire de rejeu priorisée sont fréquemment négligés dans les agents économes en échantillons, en partie à cause du coût élevé de leur intégration et du risque que l'exploration intrinsèque gaspille des budgets d'interaction rares sur des expériences sans rapport avec la tâche. Ce paysage reflète l'ère pré-Rainbow en RL profond, où des améliorations validées individuellement n'avaient pas été systématiquement combinées.

Méthodologie
Les auteurs introduisent Simulus, un agent modulaire de modèle du monde basé sur des tokens conçu pour intégrer et valider une série d'améliorations complémentaires. L'agent est construit sur l'architecture REM et comprend quatre modules distincts : un module de représentation ($V$), un modèle du monde ($M$), un contrôleur ($C$) et une mémoire de rejeu.

1. Cadre de Tokenisation Flexible ($V$) : Simulus emploie un système de tokenisation modulaire supportant des combinaisons arbitraires de modalités d'observation et d'action. Il gère nativement les images, les vecteurs continus, les variables catégorielles et les « catégoriels 2D » (grilles d'images de variables catégorielles). Les encodeurs produisent des séquences de tokens de longueur fixe, concaténées en une représentation unifiée. Cela permet à l'agent d'opérer dans des environnements divers, y compris le contrôle continu et les tâches symboliques multi-modales.
2. Modèle du Monde ($M$) : Le modèle de dynamique utilise un Réseau Rétentif (RetNet) augmenté d'un mécanisme de Prédiction Parallèle des Observations (POP) pour modéliser la dynamique de l'environnement comme une modélisation de séquence sur des tokens discrets.
   • Motivation Intrinsèque : Pour réduire l'incertitude épistémique, Simulus emploie un ensemble de têtes de prédiction de la prochaine observation. Le désaccord entre ces têtes, mesuré via la Divergence de Jensen-Shannon (JSD), génère des récompenses intrinsèques ($r_{int}$). Celles-ci sont combinées aux récompenses extrinsèques pour guider le contrôleur vers des régions à haute incertitude durant l'« imagination » (planification au sein du modèle).
   • Mémoire de Rejeu Priorisée : La mémoire de rejeu conserve les valeurs de perte du modèle du monde pour les exemples stockés. L'entraînement utilise un mélange d'échantillonnage uniforme et priorisé (contrôlé par un seul paramètre $\alpha$), où les échantillons sont sélectionnés proportionnellement à leurs pertes transformées par softmax.
3. Contrôleur ($C$) : Le contrôleur est une architecture acteur-critic étendue qui opère entièrement sur l'expérience simulée générée par le modèle du monde.
   • Régression-Comme-Classification (RaC) : Pour gérer de manière robuste les signaux de récompense et de retour non bornés, Simulus adopte un réseau de classification avec des bins exponentiels pour la prédiction de valeur. Il utilise la méthode de perte HL-Gauss (basée sur l'entropie croisée) plutôt que l'erreur quadratique moyenne traditionnelle, améliorant la stabilité de l'estimation de valeur.
4. Cycle d'Entraînement : L'agent suit un cycle de collecte de données, d'apprentissage de représentation, d'apprentissage du modèle du monde et d'apprentissage du contrôle en imagination. Le contrôleur est entraîné à l'aide de retours $\lambda$ et d'un objectif REINFORCE avec une ligne de base de valeur.

Contributions Clés

• Intégration de Composants Négligés : L'article identifie et intègre la motivation intrinsèque, la mémoire de rejeu priorisée du modèle du monde, et une approche hybride RaC. Il démontre que ces composants sont non seulement efficaces individuellement, mais produisent des gains synergiques lorsqu'ils sont combinés.
• Agents Basés sur des Tokens Versatiles : En adaptant la tokenisation multi-modale au régime économe en échantillons, Simulus établit les TBWM comme des agents versatiles capables de traiter des entrées visuelles, continues et multi-modales, dépassant les limitations des modèles d'actions discrètes uniquement basés sur l'image.
• Validation de l'Exploration Induite par le Modèle : Ce travail montre que l'exploration induite par le modèle du monde améliore significativement les performances même sous des budgets d'interaction stricts, remettant en question l'idée que la motivation intrinsèque est trop coûteuse pour les paramètres économes en échantillons.
• Fondation Open-Source : Les auteurs fournissent un point de départ solide pour la recherche sur les TBWM en open-sourçant le code et les poids du modèle, abaissant la barrière à l'entrée pour le développement d'agents de modèles du monde compétitifs.

Résultats
Simulus a été évalué sur trois benchmarks diversifiés avec des budgets d'interaction stricts :

• Atari Visuel 100K : Simulus a atteint une efficacité d'échantillons de pointe (SOTA) parmi les méthodes sans planification, surpassant des bases de référence comme DreamerV3, TWM, STORM et DIAMOND. Notamment, c'est le premier modèle du monde sans planification à atteindre des scores médians et une moyenne interquantile (IQM) au niveau humain, réalisant des performances surhumaines sur 13 des 26 jeux.
• Contrôle Continu (DMC Proprioception 500K) : Dans les tâches de contrôle continu, Simulus a constamment égalé ou légèrement surpassé DreamerV3, démontrant l'efficacité des architectures basées sur des tokens dans les domaines continus avec des vocabulaires compacts.
• Symbolique Multi-Modal (Craftax-1M) : Dans un environnement de survie 2D partiellement observable combinant des grilles symboliques et des caractéristiques continues, Simulus a surpassé les méthodes concurrentes de modèles du monde et toutes les bases de référence sans modèle (y compris RND, PPO-RNN et E3B), gérant efficacement de longues séquences de tokens.

Des études d'ablation ont confirmé que chaque composant contribue individuellement, les récompenses intrinsèques s'avérant particulièrement cruciales. La combinaison de tous les composants a produit les performances les plus fortes, indiquant une synergie significative.

Signification
L'article soutient que le potentiel inexploité des agents de modèles du monde découle de la difficulté de combiner des améliorations modulaires. En s'inspirant du cadre Rainbow en DQN, Simulus démontre qu'une combinaison systématique de flexibilité de tokenisation, de motivation intrinsèque, de mémoire de rejeu priorisée et de techniques avancées de prédiction de valeur conduit à une efficacité d'échantillons de pointe. Ce travail établit que les modèles du monde basés sur des tokens constituent une approche versatile et évolutive pour divers domaines de RL, à condition que les composants soient soigneusement intégrés pour éviter les interférences souvent observées dans les architectures monolithiques. Les auteurs maintiennent que, bien que les agents basés sur des tokens fassent actuellement face à des limitations de vitesse d'entraînement par rapport à certaines bases de référence, leur conception modulaire offre des avantages distincts pour l'inférence de politique et l'évolutivité de la recherche future.