Mejora de los Modelos del Mundo Transformer para el Aprendizaje por Refuerzo con Eficiencia de Datos

Improving Transformer World Models for Data-Efficient RL

February 3, 2025

Autores: Antoine Dedieu, Joseph Ortiz, Xinghua Lou, Carter Wendelken, Wolfgang Lehrach, J Swaroop Guntupalli, Miguel Lazaro-Gredilla, Kevin Patrick Murphy

cs.AI

Resumen

Presentamos un enfoque de RL basado en modelos que logra un nuevo rendimiento de vanguardia en el desafiante benchmark Craftax-classic, un juego de supervivencia en 2D de mundo abierto que requiere que los agentes muestren una amplia gama de habilidades generales, como una fuerte generalización, exploración profunda y razonamiento a largo plazo. Con una serie de decisiones de diseño cuidadosas destinadas a mejorar la eficiencia de muestra, nuestro algoritmo de MBRL logra una recompensa del 67.4% después de solo 1 millón de pasos de entorno, superando significativamente a DreamerV3, que alcanza el 53.2%, y, por primera vez, supera el rendimiento humano del 65.0%. Nuestro método comienza construyendo una línea de base libre de modelos de vanguardia, utilizando una arquitectura de política novedosa que combina CNNs y RNNs. Luego agregamos tres mejoras al entorno MBRL estándar: (a) "Dyna con calentamiento", que entrena la política con datos reales e imaginarios, (b) "tokenizador de vecino más cercano" en parches de imagen, que mejora el esquema para crear las entradas del modelo del mundo transformador (TWM), y (c) "forzamiento de maestro de bloque", que permite al TWM razonar conjuntamente sobre los tokens futuros del próximo paso de tiempo.

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.