コンテキストを考慮したトークン化による効率的なワールドモデル

Efficient World Models with Context-Aware Tokenization

June 27, 2024

著者: Vincent Micheli, Eloi Alonso, François Fleuret

cs.AI

要旨

深層強化学習（RL）手法のスケールアップは大きな課題となっています。生成モデリングの進展に続き、モデルベースRLは有力な候補として位置づけられています。最近のシーケンスモデリングの進歩により、効果的なトランスフォーマーベースの世界モデルが開発されましたが、環境を正確にシミュレートするために必要な長いトークンシーケンスのため、重い計算コストが伴います。本研究では、Delta-IRISという新しいエージェントを提案します。このエージェントの世界モデルアーキテクチャは、時間ステップ間の確率的デルタをエンコードする離散オートエンコーダと、連続トークンを用いて現在の世界の状態を要約し未来のデルタを予測する自己回帰型トランスフォーマーで構成されています。Crafterベンチマークにおいて、Delta-IRISは複数のフレーム予算で新たな最高性能を達成し、従来のアテンションベースのアプローチよりも一桁高速に学習することができました。私たちはコードとモデルを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.