Improving Transformer World Models for Data-Efficient RL (2502.01591v2)

Abstract: We present an approach to model-based RL that achieves a new state of the art performance on the challenging Craftax-classic benchmark, an open-world 2D survival game that requires agents to exhibit a wide range of general abilities -- such as strong generalization, deep exploration, and long-term reasoning. With a series of careful design choices aimed at improving sample efficiency, our MBRL algorithm achieves a reward of 69.66% after only 1M environment steps, significantly outperforming DreamerV3, which achieves 53.2%, and, for the first time, exceeds human performance of 65.0%. Our method starts by constructing a SOTA model-free baseline, using a novel policy architecture that combines CNNs and RNNs. We then add three improvements to the standard MBRL setup: (a) "Dyna with warmup", which trains the policy on real and imaginary data, (b) "nearest neighbor tokenizer" on image patches, which improves the scheme to create the transformer world model (TWM) inputs, and (c) "block teacher forcing", which allows the TWM to reason jointly about the future tokens of the next timestep.

• The paper introduces a novel MBRL algorithm that significantly improves sample efficiency on the Craftax-classic benchmark by achieving a reward of 67.42% after 1M steps.
• It employs a unique policy architecture combining CNNs and RNNs with enhancements like Dyna with Warmup, nearest neighbor tokenization, and block teacher forcing.
• Empirical results demonstrate that the approach outperforms previous methods and human performance, highlighting its impact on data-efficient reinforcement learning.

Improving Transformer World Models for Data-Efficient RL

The research paper under review presents a new model-based reinforcement learning (MBRL) approach that advances the state-of-the-art performance on Craftax-classic, a benchmark environment for evaluating RL agents. The Craftax-classic presents unique challenges due to its open-world 2D survival game design, requiring strong generalization, exploration, and long-term reasoning abilities from agents. Typically, the methods proposed by the authors focus on improving sample efficiency, a crucial aspect of RL that aims to reduce the required interactions with the environment to achieve competent performance.

Contributions of the Work

The authors introduce an MBRL algorithm that achieves significant improvements in performance metrics. Specifically, their method attains a reward of 67.42% after 1 million environment steps, surpassing the performance of previous methods like DreamerV3, which achieves 53.2%, and even exceeding human performance, which is rated at 65.0%.

The core contributions of the paper can be summarized as follows:

1. Model-Free Baseline Building: Initially, the authors construct an optimal model-free baseline, employing a unique policy architecture that combines Convolutional Neural Networks (CNNs) with Recurrent Neural Networks (RNNs). This architecture forms the foundational layer of their research into model-based methods.
2. MBRL Enhancements: They propose three distinct enhancements to conventional MBRL setups:
   • Dyna with Warmup: This technique involves training the policy on a mix of real and simulated data. Such an approach leverages the strengths of both model-based and model-free RL strategies.
   • Nearest Neighbor Tokenizer: To improve transformer world model (TWM) inputs, the authors develop a nearest neighbor tokenizer for discretizing image patches. This method aims to enhance data representation for TWMs.
   • Block Teacher Forcing: This method allows the TWM to consider future tokens' interdependencies within a timeline, enabling joint reasoning over the sequence data.
3. Empirical Evidence: The extensive experimental results provided indicate that these combined methodologies lead to superior performance, significantly surpassing prior models like IRIS and DreamerV2/V3 in terms of reward scores on the Craftax-classic benchmark.

Practical and Theoretical Implications

From a practical standpoint, the proposed approach contributes to creating RL models capable of operating with lower data requirements. This is crucial for applications where acquiring data is expensive or time-consuming. The improvements could be particularly beneficial for tasks requiring robust policy learning in environments with extensive state spaces and complex dynamics.

On a theoretical level, the integration of model-free and model-based components into a cohesive framework as displayed in "Dyna with Warmup" presents an interesting direction for future work. This approach could lead to novel insights into the discipline of reinforcement learning, specifically in how imaginary trajectories and memory-based policies can be best utilized to converge on optimal strategies.

Future Developments and Speculation

The paper's methodological advancements suggest several potential avenues for further research and development. One might investigate the potential of adaptive tokenization techniques that could dynamically adjust representation granularity based on context or even transfer these methodologies to other domains in artificial intelligence, such as natural language processing or autonomous systems.

Finally, the contrast and harmony between MFRL and MBRL approaches merit deeper exploration, particularly how these frameworks could be generalized across varied RL environments beyond simulated 2D survival games. The efficacy demonstrated in surpassing human-level performance also raises intriguing questions about the potential for such systems to solve real-world problems that have traditionally required human ingenuity and perception.