Voorbij A*: Betere Planning met Transformers via Zoekdynamiek Bootstrapping

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

February 21, 2024

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

cs.AI

Samenvatting

Hoewel Transformers enorme vooruitgang hebben mogelijk gemaakt in verschillende toepassingsgebieden, blijven dergelijke architecturen achter bij traditionele symbolische planners bij het oplossen van complexe besluitvormingstaken. In dit werk demonstreren we hoe Transformers kunnen worden getraind om complexe planningsproblemen op te lossen en presenteren we Searchformer, een Transformermodel dat voorheen onbekende Sokoban-puzzels in 93,7% van de gevallen optimaal oplost, terwijl het tot 26,8% minder zoekstappen gebruikt dan standaard A*-zoeken. Searchformer is een encoder-decoder Transformermodel dat is getraind om de zoekdynamiek van A* te voorspellen. Dit model wordt vervolgens verfijnd via expertiteraties om minder zoekstappen te gebruiken dan A*-zoeken, terwijl het nog steeds een optimaal plan genereert. In onze trainingsmethode wordt de zoekdynamiek van A* uitgedrukt als een tokenreeks die beschrijft wanneer taaktoestanden worden toegevoegd en verwijderd uit de zoekboom tijdens symbolisch plannen. In onze ablatiestudies over navigatie in doolhoven constateren we dat Searchformer aanzienlijk beter presteert dan baseline-modellen die het optimale plan direct voorspellen, met een 5-10 keer kleiner model en een 10 keer kleinere trainingsdataset. We demonstreren ook hoe Searchformer schaalt naar grotere en complexere besluitvormingstaken zoals Sokoban, met een verbeterd percentage opgeloste taken en verkorte zoekdynamiek.

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.