"Improving Transformer World Models for Data-Efficient RL"

Below podcast on this paper is generated with Google's Illuminate.

https://arxiv.org/abs/2502.01591

The paper addresses the challenge of creating sample-efficient reinforcement learning agents for complex open-world games requiring diverse skills. Current methods struggle with generalization, exploration, and long-term reasoning in such environments with limited data.

This paper introduces a model-based reinforcement learning method. It combines Dyna-style training with a novel tokenizer and transformer training approach to enhance sample efficiency and surpass human performance.

-----

📌 Dyna with warmup strategically blends environment interaction with model-based rollouts. This hybrid approach stabilizes policy learning. It prevents divergence from inaccurate world model predictions.

📌 Nearest neighbor tokenizer provides a stationary token space. This simplifies world model learning by grounding visual representations. The online codebook adaptation allows for open-ended environment interaction.

📌 Block teacher forcing accelerates Transformer World Model training. Parallel token prediction within timesteps enhances temporal reasoning. It mitigates error accumulation typical in autoregressive models.

----------

Methods Explored in this Paper 🔧:

→ The paper starts with a strong model-free reinforcement learning baseline. This baseline uses a new policy architecture combining Convolutional Neural Networks and Recurrent Neural Networks.

→ To enhance this baseline, the paper incorporates model-based reinforcement learning techniques. A core component is "Dyna with warmup". This method trains the policy using both real environment data and data generated by a Transformer World Model. This hybrid approach boosts sample efficiency.

→ The paper introduces a "nearest neighbor tokenizer" for processing visual inputs. This tokenizer operates on image patches. It builds a codebook of patches online. When encoding a patch, it finds the nearest code in the codebook. If no code is close enough based on a Euclidean distance threshold, a new code is added. This creates a stationary codebook for the Transformer World Model.

→ "Block teacher forcing" is proposed for training the Transformer World Model. Unlike standard autoregressive training, block teacher forcing predicts all future tokens within a timestep in parallel. This allows the model to jointly reason about future states and speeds up training and generation.

-----

Key Insights 💡:

→ Dyna with warmup effectively leverages both real and imagined data. This significantly improves policy learning and sample efficiency in model-based reinforcement learning.

→ Nearest neighbor tokenizer, when applied to image patches, creates a stable and efficient visual tokenization. This stationary codebook simplifies world model learning and improves rollout quality.

→ Block teacher forcing enhances Transformer World Model training. It accelerates training and inference. It also improves the accuracy of the world model by enabling parallel token prediction within a timestep.

-----

Results 📊:

→ Achieves 67.42% reward on Craftax-classic after 1 million environment steps.

→ Outperforms DreamerV3 which achieves 53.2% reward on the same benchmark.

→ Exceeds human expert performance level of 65.0% reward in Craftax-classic.