Skip to yearly menu bar Skip to main content


Poster

Efficient World Models with Time-Aware and Context-Augmented Tokenization

Vincent Micheli · Eloi Alonso · François Fleuret


Abstract: Scaling up deep Reinforcement Learning (RL) methods beyond traditional benchmarks presents a significant challenge. Following developments in generative modelling, model-based RL positions itself as a strong contender to bring autonomous agents to new heights.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.Herein, 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 particular, $\Delta$-IRIS sets a new state of the art at multiple frame budgets in the Crafter benchmark, while being an order of magnitude faster to train than previous attention-based approaches. We release our code and models at x.

Live content is unavailable. Log in and register to view live content