Efficient World Models with Context-Aware Tokenization

Vincent Micheli; Eloi Alonso; François Fleuret

返回文献探索

Paper · arXiv 2406.19320

Efficient World Models with Context-Aware Tokenization

Vincent Micheli, Eloi Alonso, François Fleuret

8 upvotesJune 27, 2024arXiv 预印本

AI 摘要

A new agent, $\Delta$-IRIS, achieves state-of-the-art performance in model-based reinforcement learning by using a discrete autoencoder and autoregressive transformer architecture, resulting in faster training and improved simulation efficiency.

Reinforcement Learning (RL)model-based RLsequence modellingtransformer-based world modelsdiscrete autoencoderautoregressive transformerCrafter benchmark

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.