Accurate and Efficient World Modeling with Masked Latent Transformers
This work addresses a key limitation in world modeling for reinforcement learning agents, offering a more efficient and accurate alternative to existing methods, though it is incremental in nature.
The paper tackled the problem of information loss in compressed latent spaces of world models, which negatively affects agent performance, by proposing EMERALD, an efficient and accurate world model that achieved new state-of-the-art performance on the Crafter benchmark, surpassing human experts within 10M steps and unlocking all 22 achievements.
The Dreamer algorithm has recently obtained remarkable performance across diverse environment domains by training powerful agents with simulated trajectories. However, the compressed nature of its world model's latent space can result in the loss of crucial information, negatively affecting the agent's performance. Recent approaches, such as $Δ$-IRIS and DIAMOND, address this limitation by training more accurate world models. However, these methods require training agents directly from pixels, which reduces training efficiency and prevents the agent from benefiting from the inner representations learned by the world model. In this work, we propose an alternative approach to world modeling that is both accurate and efficient. We introduce EMERALD (Efficient MaskEd latent tRAnsformer worLD model), a world model using a spatial latent state with MaskGIT predictions to generate accurate trajectories in latent space and improve the agent performance. On the Crafter benchmark, EMERALD achieves new state-of-the-art performance, becoming the first method to surpass human experts performance within 10M environment steps. Our method also succeeds to unlock all 22 Crafter achievements at least once during evaluation.