Effiziente Weltmodelle mit kontextsensitiver Tokenisierung.

Efficient World Models with Context-Aware Tokenization

June 27, 2024

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

cs.AI

Zusammenfassung

Die Skalierung von Methoden des tiefen Reinforcement-Lernens (RL) stellt eine bedeutende Herausforderung dar. Im Anschluss an Entwicklungen im Bereich der generativen Modellierung positioniert sich modellbasiertes RL als ein starker Konkurrent. Jüngste Fortschritte in der Sequenzmodellierung haben zu effektiven weltbasierten Transformer-Modellen geführt, wenn auch zu dem Preis schwerer Berechnungen aufgrund langer Sequenzen von Tokens, die erforderlich sind, um Umgebungen genau zu simulieren. In dieser Arbeit schlagen wir Delta-IRIS vor, einen neuen Agenten mit einer Weltmodellarchitektur, die aus einem diskreten Autoencoder besteht, der stochastische Deltas zwischen Zeitpunkten codiert, und einem autoregressiven Transformer, der zukünftige Deltas vorhersagt, indem er den aktuellen Zustand der Welt mit kontinuierlichen Tokens zusammenfasst. Im Crafter-Benchmark setzt Delta-IRIS einen neuen Stand der Technik bei mehreren Frame-Budgets, während es um eine Größenordnung schneller zu trainieren ist als frühere auf Aufmerksamkeit basierende Ansätze. Wir veröffentlichen unseren Code und Modelle unter 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.