Au-delà d'A* : Une meilleure planification avec les Transformers via la dynamique de recherche 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

Résumé

Bien que les Transformers aient permis des progrès considérables dans divers contextes d'application, ces architectures restent encore en retard par rapport aux planificateurs symboliques traditionnels pour résoudre des tâches complexes de prise de décision. Dans ce travail, nous démontrons comment entraîner des Transformers à résoudre des tâches de planification complexes et présentons Searchformer, un modèle Transformer qui résout de manière optimale des puzzles Sokoban jamais vus auparavant dans 93,7 % des cas, tout en utilisant jusqu'à 26,8 % d'étapes de recherche en moins que la recherche A* standard. Searchformer est un modèle Transformer encodeur-décodeur entraîné à prédire la dynamique de recherche de A*. Ce modèle est ensuite affiné via des itérations expertes pour effectuer moins d'étapes de recherche que la recherche A* tout en générant un plan optimal. Dans notre méthode d'entraînement, la dynamique de recherche de A* est exprimée sous forme d'une séquence de tokens décrivant quand les états de la tâche sont ajoutés et retirés de l'arbre de recherche pendant la planification symbolique. Dans nos études d'ablation sur la navigation dans les labyrinthes, nous constatons que Searchformer surpasse significativement les modèles de référence qui prédisent directement le plan optimal, avec une taille de modèle 5 à 10 fois plus petite et un ensemble de données d'entraînement 10 fois plus réduit. Nous démontrons également comment Searchformer s'adapte à des tâches de prise de décision plus vastes et complexes comme Sokoban, avec un pourcentage amélioré de tâches résolues et une dynamique de recherche raccourcie.

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.