Efficiënte Wereldmodellen met Contextbewuste Tokenisatie

Efficient World Models with Context-Aware Tokenization

June 27, 2024

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

cs.AI

Samenvatting

Het opschalen van deep Reinforcement Learning (RL)-methoden vormt een aanzienlijke uitdaging. In navolging van ontwikkelingen in generatieve modellering positioneert modelgebaseerde RL zich als een sterke kandidaat. Recente vooruitgang in sequentiemodellering heeft geleid tot effectieve transformer-gebaseerde wereldmodellen, zij het tegen de prijs van zware berekeningen vanwege de lange sequenties van tokens die nodig zijn om omgevingen nauwkeurig te simuleren. In dit werk stellen we Delta-IRIS voor, een nieuwe agent met een wereldmodelarchitectuur die bestaat uit een discrete auto-encoder die stochastische delta's tussen tijdstappen codeert en een autoregressieve transformer die toekomstige delta's voorspelt door de huidige staat van de wereld samen te vatten met continue tokens. In de Crafter-benchmark vestigt Delta-IRIS een nieuwe standaard op meerdere frame-budgetten, terwijl het een orde van grootte sneller is om te trainen dan eerdere aandacht-gebaseerde benaderingen. We maken onze code en modellen beschikbaar op 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.