Jenseits von A*: Bessere Planung mit Transformern durch Suchdynamik Bootstrapping

Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping

February 21, 2024

Autoren: Lucas Lehnert, Sainbayar Sukhbaatar, Paul Mcvay, Michael Rabbat, Yuandong Tian

cs.AI

Zusammenfassung

Während Transformer enorme Fortschritte in verschiedenen Anwendungsbereichen ermöglicht haben, hinken solche Architekturen bei der Lösung komplexer Entscheidungsfindungsaufgaben immer noch traditionellen symbolischen Planern hinterher. In dieser Arbeit zeigen wir, wie man Transformer trainieren kann, um komplexe Planungsaufgaben zu lösen, und präsentieren Searchformer, ein Transformer-Modell, das zuvor unbekannte Sokoban-Rätsel in 93,7 % der Fälle optimal löst und dabei bis zu 26,8 % weniger Suchschritte als die Standard-A^*-Suche verwendet. Searchformer ist ein Encoder-Decoder-Transformer-Modell, das darauf trainiert ist, die Suchdynamik von A^* vorherzusagen. Dieses Modell wird dann durch Expert-Iterationen feinabgestimmt, um weniger Suchschritte als die A^*-Suche durchzuführen und dennoch einen optimalen Plan zu generieren. In unserer Trainingsmethode wird die Suchdynamik von A^* als eine Token-Sequenz dargestellt, die beschreibt, wann Aufgabenstatus während der symbolischen Planung in den Suchbaum hinzugefügt und entfernt werden. In unseren Ablationsstudien zur Navigation in Labyrinthen stellen wir fest, dass Searchformer Baseline-Modelle, die den optimalen Plan direkt vorhersagen, deutlich übertrifft, und dies mit einer 5-10 Mal kleineren Modellgröße und einem 10 Mal kleineren Trainingsdatensatz. Wir zeigen auch, wie Searchformer auf größere und komplexere Entscheidungsfindungsaufgaben wie Sokoban skaliert, mit einer verbesserten Lösungshäufigkeit und verkürzter Suchdynamik.

English

While Transformers have enabled tremendous progress in various application settings, such architectures still lag behind traditional symbolic planners for solving complex decision making tasks. In this work, we demonstrate how to train Transformers to solve complex planning tasks and present Searchformer, a Transformer model that optimally solves previously unseen Sokoban puzzles 93.7% of the time, while using up to 26.8% fewer search steps than standard A^* search. Searchformer is an encoder-decoder Transformer model trained to predict the search dynamics of A^*. This model is then fine-tuned via expert iterations to perform fewer search steps than A^* search while still generating an optimal plan. In our training method, A^*'s search dynamics are expressed as a token sequence outlining when task states are added and removed into the search tree during symbolic planning. In our ablation studies on maze navigation, we find that Searchformer significantly outperforms baselines that predict the optimal plan directly with a 5-10times smaller model size and a 10times smaller training dataset. We also demonstrate how Searchformer scales to larger and more complex decision making tasks like Sokoban with improved percentage of solved tasks and shortened search dynamics.