Verbesserung von Transformer-Weltmodellen für dateneffizientes RL

Improving Transformer World Models for Data-Efficient RL

February 3, 2025

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

cs.AI

Zusammenfassung

Wir präsentieren einen Ansatz für modellbasiertes RL, der eine neue Bestleistung auf dem anspruchsvollen Craftax-classic Benchmark erzielt, einem Open-World-2D-Überlebensspiel, bei dem Agenten eine Vielzahl von allgemeinen Fähigkeiten zeigen müssen - wie starke Verallgemeinerung, tiefe Exploration und langfristiges Denken. Mit einer Reihe sorgfältiger Designentscheidungen zur Verbesserung der Stichproben-Effizienz erreicht unser MBRL-Algorithmus eine Belohnung von 67,4 % nach nur 1 Mio. Umgebungsschritten, was deutlich besser ist als DreamerV3 mit 53,2 % und erstmals die menschliche Leistung von 65,0 % übertrifft. Unser Ansatz beginnt mit dem Aufbau einer SOTA modellfreien Baseline, die eine neuartige Richtlinienarchitektur verwendet, die CNNs und RNNs kombiniert. Anschließend fügen wir drei Verbesserungen zum Standard-MBRL-Setup hinzu: (a) "Dyna mit Warm-up", das die Richtlinie auf echten und imaginären Daten trainiert, (b) "nächster Nachbar-Tokenizer" auf Bildausschnitten, der das Schema zur Erstellung der Eingaben des Transformer-Weltmodells (TWM) verbessert, und (c) "Block-Lehrerzwang", der es dem TWM ermöglicht, gemeinsam über die zukünftigen Token des nächsten Zeitschritts zu spekulieren.

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.