Amélioration des modèles de monde Transformer pour un RL efficace en données

Improving Transformer World Models for Data-Efficient RL

February 3, 2025

Auteurs: Antoine Dedieu, Joseph Ortiz, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, J Swaroop Guntupalli, Miguel Lazaro-Gredilla, Kevin Patrick Murphy

cs.AI

Résumé

Nous présentons une approche de l'apprentissage par renforcement basée sur le modèle qui atteint une nouvelle performance de pointe sur le difficile benchmark Craftax-classic, un jeu de survie en monde ouvert en 2D qui demande aux agents de démontrer un large éventail d'aptitudes générales telles que la forte généralisation, l'exploration approfondie et le raisonnement à long terme. Grâce à une série de choix de conception minutieux visant à améliorer l'efficacité des échantillons, notre algorithme MBRL atteint une récompense de 67,4 % après seulement 1 million d'étapes dans l'environnement, surpassant significativement DreamerV3, qui atteint 53,2 %, et, pour la première fois, dépasse la performance humaine de 65,0 %. Notre méthode commence par construire une ligne de base sans modèle de pointe, en utilisant une architecture de politique novatrice combinant des CNN et des RNN. Nous ajoutons ensuite trois améliorations à la configuration MBRL standard : (a) "Dyna avec préchauffage", qui entraîne la politique sur des données réelles et imaginaires, (b) "tokeniseur de voisin le plus proche" sur des fragments d'image, améliorant le schéma pour créer les entrées du modèle mondial transformateur (TWM), et (c) "forcing de l'enseignant par bloc", permettant au TWM de raisonner conjointement sur les jetons futurs de l'étape suivante.

English

We present an approach to model-based RL that achieves a new state of the art performance on the challenging Craftax-classic benchmark, an open-world 2D survival game that requires agents to exhibit a wide range of general abilities -- such as strong generalization, deep exploration, and long-term reasoning. With a series of careful design choices aimed at improving sample efficiency, our MBRL algorithm achieves a reward of 67.4% after only 1M environment steps, significantly outperforming DreamerV3, which achieves 53.2%, and, for the first time, exceeds human performance of 65.0%. Our method starts by constructing a SOTA model-free baseline, using a novel policy architecture that combines CNNs and RNNs. We then add three improvements to the standard MBRL setup: (a) "Dyna with warmup", which trains the policy on real and imaginary data, (b) "nearest neighbor tokenizer" on image patches, which improves the scheme to create the transformer world model (TWM) inputs, and (c) "block teacher forcing", which allows the TWM to reason jointly about the future tokens of the next timestep.