Улучшение моделей миров трансформера для эффективного обучения с подкреплением на данных.

Improving Transformer World Models for Data-Efficient RL

February 3, 2025

Авторы: Antoine Dedieu, Joseph Ortiz, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, J Swaroop Guntupalli, Miguel Lazaro-Gredilla, Kevin Patrick Murphy

cs.AI

Аннотация

Мы представляем подход к моделированию на основе обучения с подкреплением, который достигает нового state of the art производительности на сложном бенчмарке Craftax-classic, открытой 2D игре выживания, требующей от агентов демонстрации широкого спектра общих способностей - таких как сильная обобщающая способность, глубокое исследование и долгосрочное мышление. С помощью серии тщательно продуманных выборов дизайна, направленных на улучшение эффективности использования выборок, наш алгоритм MBRL достигает награды в 67.4% после всего лишь 1 миллиона шагов в среде, значительно превосходя DreamerV3, который достигает 53.2%, и впервые превосходит человеческую производительность в 65.0%. Наш метод начинается с построения SOTA модельного базиса, используя новую архитектуру политики, объединяющую CNN и RNN. Затем мы добавляем три улучшения к стандартной настройке MBRL: (a) "Dyna с разогревом", который обучает политику на реальных и мнимых данных, (b) "токенизатор ближайшего соседа" на изображениях, улучшающий схему создания входов трансформерной модели мира (TWM), и (c) "блочное принуждение учителя", позволяющее TWM совместно рассуждать о будущих токенах следующего временного шага.

English

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 67.4% 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.