Het verbeteren van Transformer-wereldmodellen voor gegevensefficiënte RL.

Improving Transformer World Models for Data-Efficient RL

February 3, 2025

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

cs.AI

Samenvatting

We presenteren een benadering voor modelgebaseerd RL die een nieuwe state-of-the-art prestatie behaalt op de uitdagende Craftax-classic benchmark, een open-wereld 2D survival game waarbij agenten een breed scala aan algemene vaardigheden moeten vertonen - zoals sterke generalisatie, diepgaande exploratie en langetermijnredenering. Met een reeks zorgvuldige ontwerpkeuzes gericht op het verbeteren van de steekproefeffectiviteit, behaalt ons MBRL-algoritme een beloning van 67,4% na slechts 1 miljoen omgevingsstappen, aanzienlijk beter dan DreamerV3, dat 53,2% behaalt, en, voor de eerste keer, de menselijke prestatie van 65,0% overtreft. Onze methode begint met het construeren van een state-of-the-art modelvrije basislijn, met behulp van een nieuw beleidsarchitectuur die CNN's en RNN's combineert. Vervolgens voegen we drie verbeteringen toe aan de standaard MBRL-opstelling: (a) "Dyna met opstart", dat het beleid traint op echte en denkbeeldige gegevens, (b) "nearest neighbor tokenizer" op beeldpatches, dat het schema verbetert om de transformer wereldmodel (TWM) invoeren te creëren, en (c) "block teacher forcing", dat de TWM in staat stelt om gezamenlijk te redeneren over de toekomstige tokens van de volgende tijdstap.

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.