Modèles de monde efficaces avec tokenisation contextuelle

Efficient World Models with Context-Aware Tokenization

June 27, 2024

Auteurs: Vincent Micheli, Eloi Alonso, François Fleuret

cs.AI

Résumé

La mise à l'échelle des méthodes d'apprentissage par renforcement profond (RL) représente un défi majeur. À la suite des avancées dans le domaine de la modélisation générative, l'approche basée sur des modèles (model-based RL) s'impose comme un candidat sérieux. Les récents progrès en modélisation séquentielle ont conduit à des modèles du monde efficaces basés sur des transformateurs, bien qu'au prix de calculs intensifs en raison des longues séquences de tokens nécessaires pour simuler précisément les environnements. Dans ce travail, nous proposons Delta-IRIS, un nouvel agent doté d'une architecture de modèle du monde composée d'un autoencodeur discret qui encode des deltas stochastiques entre les pas de temps et d'un transformateur autorégressif qui prédit les deltas futurs en résumant l'état actuel du monde avec des tokens continus. Sur le benchmark Crafter, Delta-IRIS établit un nouvel état de l'art pour plusieurs budgets d'images, tout en étant un ordre de grandeur plus rapide à entraîner que les approches précédentes basées sur l'attention. Nous mettons à disposition notre code et nos modèles à l'adresse https://github.com/vmicheli/delta-iris.

English

Scaling up deep Reinforcement Learning (RL) methods presents a significant challenge. Following developments in generative modelling, model-based RL positions itself as a strong contender. Recent advances in sequence modelling have led to effective transformer-based world models, albeit at the price of heavy computations due to the long sequences of tokens required to accurately simulate environments. In this work, we propose Delta-IRIS, a new agent with a world model architecture composed of a discrete autoencoder that encodes stochastic deltas between time steps and an autoregressive transformer that predicts future deltas by summarizing the current state of the world with continuous tokens. In the Crafter benchmark, Delta-IRIS sets a new state of the art at multiple frame budgets, while being an order of magnitude faster to train than previous attention-based approaches. We release our code and models at https://github.com/vmicheli/delta-iris.