Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866909183515820032 |
|---|---|
| author | Lehnert, Lucas Sukhbaatar, Sainbayar Su, DiJia Zheng, Qinqing Mcvay, Paul Rabbat, Michael Tian, Yuandong |
| author_facet | Lehnert, Lucas Sukhbaatar, Sainbayar Su, DiJia Zheng, Qinqing Mcvay, Paul Rabbat, Michael Tian, Yuandong |
| contents | While Transformers have enabled tremendous progress in various application settings, such architectures still trail behind traditional symbolic planners for solving complex decision making tasks. In this work, we demonstrate how to train Transformers to solve complex planning tasks. This is accomplished by training an encoder-decoder Transformer model to predict the search dynamics of the $A^*$ search algorithm. 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 the $A^*$ implementation that was used for training initially. 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. 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2402_14083 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping Lehnert, Lucas Sukhbaatar, Sainbayar Su, DiJia Zheng, Qinqing Mcvay, Paul Rabbat, Michael Tian, Yuandong Artificial Intelligence While Transformers have enabled tremendous progress in various application settings, such architectures still trail behind traditional symbolic planners for solving complex decision making tasks. In this work, we demonstrate how to train Transformers to solve complex planning tasks. This is accomplished by training an encoder-decoder Transformer model to predict the search dynamics of the $A^*$ search algorithm. 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 the $A^*$ implementation that was used for training initially. 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. 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 |
| topic | Artificial Intelligence |
| url | https://arxiv.org/abs/2402.14083 |