Modelos de Mundo Eficientes com Tokenização Consciente do Contexto

Efficient World Models with Context-Aware Tokenization

June 27, 2024

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

cs.AI

Resumo

Aumentar os métodos de Aprendizado por Reforço Profundo (RL) apresenta um desafio significativo. Seguindo os desenvolvimentos em modelagem generativa, o RL baseado em modelo posiciona-se como um forte concorrente. Avanços recentes em modelagem de sequência levaram a modelos de mundo baseados em transformadores eficazes, embora ao custo de intensas computações devido às longas sequências de tokens necessárias para simular ambientes com precisão. Neste trabalho, propomos o Delta-IRIS, um novo agente com uma arquitetura de modelo de mundo composta por um autoencoder discreto que codifica deltas estocásticos entre passos de tempo e um transformador autoregressivo que prevê futuros deltas resumindo o estado atual do mundo com tokens contínuos. No benchmark Crafter, o Delta-IRIS estabelece um novo estado da arte em vários orçamentos de quadros, sendo uma ordem de magnitude mais rápido para treinar do que abordagens anteriores baseadas em atenção. Disponibilizamos nosso código e modelos em 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.