Além do A*: Melhor Planejamento com Transformers via Dinâmica de Busca Bootstrapping

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

February 21, 2024

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

cs.AI

Resumo

Embora os Transformers tenham possibilitado avanços significativos em diversos cenários de aplicação, essas arquiteturas ainda ficam atrás dos planejadores simbólicos tradicionais na resolução de tarefas complexas de tomada de decisão. Neste trabalho, demonstramos como treinar Transformers para resolver tarefas complexas de planejamento e apresentamos o Searchformer, um modelo Transformer que resolve de forma ótima quebra-cabeças de Sokoban nunca vistos anteriormente em 93,7% dos casos, utilizando até 26,8% menos etapas de busca do que o algoritmo A^* padrão. O Searchformer é um modelo Transformer de codificador-decodificador treinado para prever a dinâmica de busca do A^*. Esse modelo é então ajustado por meio de iterações especializadas para realizar menos etapas de busca do que o A^*, ainda assim gerando um plano ótimo. Em nosso método de treinamento, a dinâmica de busca do A^* é expressa como uma sequência de tokens que descreve quando os estados da tarefa são adicionados e removidos da árvore de busca durante o planejamento simbólico. Em nossos estudos de ablação sobre navegação em labirintos, descobrimos que o Searchformer supera significativamente as baselines que preveem o plano ótimo diretamente, com um tamanho de modelo 5 a 10 vezes menor e um conjunto de treinamento 10 vezes menor. Também demonstramos como o Searchformer escala para tarefas de tomada de decisão maiores e mais complexas, como o Sokoban, com uma porcentagem melhor de tarefas resolvidas e dinâmicas de busca reduzidas.

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.