bibtext: @inproceedings{
  lehnert2024beyond,
  title={Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping},
  author={Lucas Lehnert and Sainbayar Sukhbaatar and DiJia Su and Qinqing Zheng and Paul McVay and Michael Rabbat and Yuandong Tian},
  booktitle={First Conference on Language Modeling},
  year={2024},
  url={https://openreview.net/forum?id=SGoVIC0u0f}
  }
abstract: While Transformers have enabled tremendous progress in various application settings, such architectures still struggle with solving planning and sequential decision-making tasks. In this work, we demonstrate how to train Transformers to solve complex planning tasks. This is accomplished by first designing a synthetic language that captures the computation performed by the $A^*$ search algorithm when solving a planning task. Then, an encoder-decoder Transformer model is trained to predict this language, resulting in a language model that can correctly solve novel planning tasks by generating $A^*$'s search dynamics. We fine tune this model to obtain a 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 our $A^*$ reference implementation. Searchformer significantly outperforms baselines that predict the optimal plan directly with a 5-10$\times$ smaller model size and a 10$\times$ smaller training dataset. Lastly, we demonstrate how Searchformer scales to larger and more complex decision making tasks with improved percentage of solved tasks and shortened search dynamics.
title: Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping
authors: Lucas Lehnert Sainbayar Sukhbaatar DiJia Su Qinqing Zheng Paul McVay Michael Rabbat Yuandong Tian