Oltre A*: Migliorare la Pianificazione con i Transformer attraverso la Dinamica della Ricerca Bootstrapping

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

February 21, 2024

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

cs.AI

Abstract

Sebbene i Transformer abbiano permesso progressi significativi in vari contesti applicativi, queste architetture rimangono ancora indietro rispetto ai pianificatori simbolici tradizionali per la risoluzione di compiti decisionali complessi. In questo lavoro, dimostriamo come addestrare i Transformer a risolvere compiti di pianificazione complessi e presentiamo Searchformer, un modello Transformer che risolve in modo ottimale puzzle di Sokoban mai visti prima nel 93,7% dei casi, utilizzando fino al 26,8% in meno di passaggi di ricerca rispetto alla ricerca A^* standard. Searchformer è un modello Transformer encoder-decoder addestrato a prevedere le dinamiche di ricerca di A^*. Questo modello viene poi affinato tramite iterazioni esperte per eseguire meno passaggi di ricerca rispetto alla ricerca A^* pur generando comunque un piano ottimale. Nel nostro metodo di addestramento, le dinamiche di ricerca di A^* sono espresse come una sequenza di token che descrive quando gli stati del compito vengono aggiunti e rimossi dall'albero di ricerca durante la pianificazione simbolica. Nei nostri studi di ablazione sulla navigazione in labirinto, scopriamo che Searchformer supera significativamente i baseline che prevedono direttamente il piano ottimale con una dimensione del modello 5-10 volte più piccola e un dataset di addestramento 10 volte più ridotto. Dimostriamo inoltre come Searchformer si adatti a compiti decisionali più grandi e complessi come Sokoban, migliorando la percentuale di compiti risolti e riducendo le dinamiche di ricerca.

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.