Efficient World Models with Context-Aware Tokenization (2406.19320v1)

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.

• The paper introduces a novel model-based RL agent, 'diris', that efficiently tokenizes stochastic deltas to reduce computational overhead.
• It leverages a discrete autoencoder and an autoregressive transformer with context-aware tokenization to improve reconstruction accuracy and learning speed.
• Results on the Crafter benchmark demonstrate state-of-the-art performance, achieving a tenfold faster training rate compared to competing models.

Efficient World Models with Context-Aware Tokenization

Introduction

The paper "Efficient World Models with Context-Aware Tokenization" presents a novel model-based RL agent, \diris, for tackling the challenges of scaling reinforcement learning in complex environments. The core innovation is a world model architecture comprising a discrete autoencoder and an autoregressive transformer, designed to encode stochastic deltas between time steps efficiently. This approach aims to overcome the prohibitive computational demands of previous token-based simulation methods, thereby enabling faster and more effective learning in visually complex domains like the Crafter benchmark.

Methodology

Discrete Autoencoder

The autoencoder in \diris distinguishes itself by encoding frames conditioned on previous frames and actions, focusing solely on the stochastic deltas between time steps. This non-traditional strategy reduces the number of tokens required for encoding, which is pivotal for speeding up the subsequent autoregressive modelling. The encoder utilizes a CNN architecture to process and tokenize image observations, while the decoder reconstructs frames using these tokens along with past actions and frames, effectively disentangling deterministic and stochastic components of scene dynamics.

Figure 1: Evidence of dynamics disentanglement. Shows how \diris encodes stochastic deltas between time steps with \dtoken-tokens, crucial for efficient modelling of complex environments.

Autoregressive Transformer

Enhancements to the autoregressive transformer centre on integrating \itoken-tokens derived from continuous frame embeddings to alleviate the complexity of reasoning over past \dtoken-tokens. This integration enables the dynamics model to form a more succinct representation of the current world state and predict subsequent stochastic deltas more effectively than its predecessors.

Figure 2: Trajectories imagined with (top) and without (bottom) \itoken-tokens highlight \diris' capacity to handle intricate environment dynamics with \itoken-tokens.

Experimental Results

Crafter Benchmark Performance

Experiments demonstrate \diris' efficacy in the Crafter benchmark environment, where it sets new state-of-the-art results and surpasses DreamerV3 beyond 3M frames. Notably, \diris maintains a superior learning rate—training ten times faster than iris—and shows robustness across various frame budgets. This performance affirms the practical value of less computationally demanding tokenization in complex RL environments.

Figure 3: Returns at multiple frame budgets in the Crafter benchmark, emphasizing the superior returns \diris\ achieves compared to competing models.

Quality of World Models

Qualitative experiments underscore the effectiveness of \diris's architecture. When analyzing autoencoded test frames, \diris exhibits minimal reconstruction loss even with fewer tokens relative to iris, indicating precise modeling capabilities. Furthermore, the imagined trajectories exhibit rich gameplay mechanics, illustrating the transformer model's proficiency in simulating visually challenging environments.

Figure 4: Bottom 1\% test frames autoencoded by \diris and iris, underscoring \diris' superior reconstruction accuracy.

Implications and Future Work

\diris represents a significant step forward in the design of scalable RL agents. By efficiently separating stochastic dynamics in complex environments and optimizing computation with context-aware tokenization, it paves the way for RL systems that can operate in more nuanced and visually rich simulations. Future investigations might explore dynamic token usage based on instantaneous context and integrate the latent representations into policy learning for enhanced performance.

Conclusion

The paper offers a comprehensive exploration of efficient world modeling techniques with \diris, demonstrating its capability to simulate and learn in complex environments efficiently. This work lays groundwork for further innovations in RL, especially those seeking to balance computational efficiency with scalable complex environment interaction strategies.