Modelos de Mundo Eficientes con Tokenización Consciente del Contexto

Efficient World Models with Context-Aware Tokenization

June 27, 2024

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

cs.AI

Resumen

El escalamiento de los métodos de Aprendizaje por Refuerzo Profundo (RL, por sus siglas en inglés) representa un desafío significativo. Siguiendo los avances en modelado generativo, el RL basado en modelos se posiciona como un fuerte contendiente. Los recientes progresos en modelado de secuencias han llevado a modelos del mundo efectivos basados en transformadores, aunque a costa de cálculos intensivos debido a las largas secuencias de tokens necesarias para simular con precisión los entornos. En este trabajo, proponemos Delta-IRIS, un nuevo agente con una arquitectura de modelo del mundo compuesta por un autoencoder discreto que codifica deltas estocásticos entre pasos de tiempo y un transformador autorregresivo que predice deltas futuros resumiendo el estado actual del mundo con tokens continuos. En el benchmark Crafter, Delta-IRIS establece un nuevo estado del arte en múltiples presupuestos de frames, mientras que es un orden de magnitud más rápido de entrenar que los enfoques anteriores basados en atención. Publicamos nuestro código y modelos en 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.