Modelli Mondiali Efficienti con Tokenizzazione Consapevole del Contesto

Efficient World Models with Context-Aware Tokenization

June 27, 2024

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

cs.AI

Abstract

L'adozione su larga scala dei metodi di apprendimento per rinforzo profondo (Reinforcement Learning, RL) rappresenta una sfida significativa. Seguendo gli sviluppi nella modellazione generativa, l'RL basato su modelli si posiziona come un forte contendente. I recenti progressi nella modellazione di sequenze hanno portato a modelli del mondo efficaci basati su trasformatori, sebbene al prezzo di calcoli intensivi dovuti alle lunghe sequenze di token necessarie per simulare accuratamente gli ambienti. In questo lavoro, proponiamo Delta-IRIS, un nuovo agente con un'architettura di modello del mondo composta da un autoencoder discreto che codifica delta stocastici tra i passaggi temporali e un trasformatore autoregressivo che prevede i delta futuri sintetizzando lo stato attuale del mondo con token continui. Nel benchmark Crafter, Delta-IRIS stabilisce un nuovo stato dell'arte per diversi budget di frame, risultando un ordine di grandezza più veloce da addestrare rispetto ai precedenti approcci basati sull'attenzione. Rilasciamo il nostro codice e i modelli all'indirizzo 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.