Эффективные модели мира с контекстно-ориентированной токенизацией.

Efficient World Models with Context-Aware Tokenization

June 27, 2024

Авторы: Vincent Micheli, Eloi Alonso, François Fleuret

cs.AI

Аннотация

Масштабирование методов глубокого обучения с подкреплением (RL) представляет собой значительное вызов. Вслед за разработками в области генеративного моделирования, модельно-ориентированное обучение с подкреплением позиционирует себя как сильный конкурент. Недавние достижения в моделировании последовательностей привели к эффективным моделям мира на основе трансформеров, хотя за счет тяжелых вычислений из-за длинных последовательностей токенов, необходимых для точного моделирования окружающей среды. В данной работе мы предлагаем Delta-IRIS, нового агента с архитектурой модели мира, состоящей из дискретного автоэнкодера, который кодирует стохастические дельты между временными шагами, и авторегрессионного трансформера, который предсказывает будущие дельты, суммируя текущее состояние мира с непрерывными токенами. В бенчмарке Crafter, Delta-IRIS устанавливает новый state of the art при нескольких бюджетах кадров, при этом обучается на порядок быстрее предыдущих подходов, основанных на внимании. Мы публикуем наш код и модели по адресу 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.