Efficient World Models with Context-Aware Tokenization
Authors: name: Vincent Micheli; name: Eloi Alonso; name: François Fleuret
Published: 2024-06-27 16:54:12

Abstract: 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.