Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping (2402.14083v2)

Abstract: 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.

• The paper introduces Searchformer, a novel Transformer model that bootstraps A* search dynamics to achieve superior planning performance.
• It demonstrates a 93.7% success rate in solving Sokoban puzzles while reducing search steps by 26.8% compared to traditional methods.
• The study bridges deep learning and symbolic planning, paving the way for more efficient decision-making in robotics, logistics, and complex applications.

Enhancing Transformer Efficiency in Planning Tasks through Search Dynamics Bootstrapping

Introduction to Searchformer

Recent advancements in deep learning and, particularly, Transformer architectures have markedly impacted various fields, from natural language processing to computer vision. However, these models often encounter difficulties in complex decision-making tasks when compared to traditional symbolic planners. Our focus is on a novel approach that leverages the strengths of Transformers for elaborate planning tasks, which are traditionally out of reach for such models. We introduce Searchformer, a Transformer-based model, that not only competes with but also surpasses the A* search algorithm in solving complex Sokoban puzzles with more efficiency.

Methodology

Searchformer is built upon an encoder-decoder Transformer architecture, initially trained to imitate the search dynamics of the A* search algorithm. It is subsequently fine-tuned through expert iterations, thereby improving its search efficiency beyond that of its symbolic planning counterpart. This process involves generating a synthetic dataset by simulating the A*search, where the search dynamics are recorded as token sequences. These sequences intricately describe the adding and removing of task states into a search tree. Crucially, this method allows the Searchformer to be initially versed in traditional search techniques and then transcend them by optimizing the search steps required to achieve an optimal plan.

Experimentation and Results

Our experiments were conducted across two main domains: maze navigation tasks and Sokoban puzzles. The results indicate a remarkable capability of Searchformer to outdo the efficiency of A* search. Notably, in solving unseen Sokoban puzzles, Searchformer demonstrated a 93.7% success rate while reducing the search steps by an average of 26.8%. These improvements were achieved through a methodical bootstrapping process, where the Transformer model is iteratively fine-tuned to refine search efficiency while maintaining or enhancing solution accuracy. Moreover, our ablation studies further affirm the importance of including search dynamics in training data, highlighting that models trained with search-augmented sequences drastically outperform their solution-only counterparts, particularly in low-data regimes or when tackling more complex tasks.

Theoretical and Practical Implications

This work underscores the potential of Transformer models in solving complex decision-making tasks, a domain traditionally dominated by symbolic planning algorithms. From a theoretical standpoint, it bridges the gap between deep learning and symbolic planning methodologies, proving that with adequate training, Transformers can internalize and evolve beyond established search algorithms. Practically, this research paves the way for more efficient automated planning systems, potentially benefiting a wide range of applications, from robotics to logistics and beyond. Furthermore, the efficiency gains in search dynamics could lead to significant computational cost reductions, making sophisticated planning solutions more accessible.

Future Directions and Broader Impact

While the current implementation of Searchformer marks a significant advancement, there is room for exploration. Future work could delve into curriculum learning strategies or integrate hierarchical planning methods to further enhance model efficiency and capability. Addressing these challenges could widen the applicability of AI in scenarios requiring complex decision-making under constraints.

The broader impact of this work is multifaceted, potentially influencing both academic research directions and real-world applications in industries reliant on planning and scheduling. Nevertheless, it's critical to continue assessing the ethical implications and ensuring these advancements contribute positively to society.

Acknowledgements

This section acknowledges contributions from individuals who provided insights or feedback throughout the research process, highlighting the collaborative nature of this endeavor.

Conclusion

Searchformer represents a significant step toward harnessing the power of Transformer models for complex planning tasks, showcasing the ability to learn and improve upon traditional planning algorithms. This work not only challenges the current limitations of deep learning models in strategic decision-making but also opens new avenues for AI research and application in areas where efficient planning is crucial.