Miglioramento dei modelli di mondo Transformer per RL efficiente nei dati.

Improving Transformer World Models for Data-Efficient RL

February 3, 2025

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

cs.AI

Abstract

Presentiamo un approccio all'apprendimento per rinforzo basato su modelli che raggiunge una nuova performance di stato dell'arte sulla sfida della benchmark Craftax-classic, un gioco di sopravvivenza 2D open-world che richiede agli agenti di mostrare una vasta gamma di abilità generali - come una forte generalizzazione, esplorazione approfondita e ragionamento a lungo termine. Con una serie di scelte di progettazione oculate mirate a migliorare l'efficienza del campionamento, il nostro algoritmo di apprendimento per rinforzo basato su modelli raggiunge un punteggio di 67,4% dopo soli 1M di passaggi nell'ambiente, superando significativamente DreamerV3, che raggiunge il 53,2%, e, per la prima volta, supera la performance umana del 65,0%. Il nostro metodo inizia costruendo un punto di riferimento senza modelli all'avanguardia, utilizzando un'architettura di politica innovativa che combina CNN e RNN. Aggiungiamo poi tre miglioramenti alla configurazione standard dell'apprendimento per rinforzo basato su modelli: (a) "Dyna con riscaldamento", che addestra la politica su dati reali e immaginari, (b) "tokenizzatore del vicino più prossimo" su patch di immagini, che migliora lo schema per creare gli input del modello del mondo trasformatore (TWM), e (c) "forzatura dell'insegnante a blocchi", che consente al TWM di ragionare congiuntamente sui token futuri del prossimo passaggio temporale.

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.