A*を超えて：探索ダイナミクスを活用したTransformerによる計画の改善 ブートストラップ法

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

February 21, 2024

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

cs.AI

要旨

Transformerは様々なアプリケーション分野で大きな進歩をもたらしてきたが、複雑な意思決定タスクを解決する際には、依然として従来の記号的プランナーに後れを取っている。本研究では、Transformerを訓練して複雑なプランニングタスクを解決する方法を示し、SearchformerというTransformerモデルを提案する。Searchformerは、未見のソコバンパズルを93.7%の確率で最適に解決し、標準的なA^*探索と比べて最大26.8%少ない探索ステップを使用する。Searchformerは、A^*の探索ダイナミクスを予測するように訓練されたエンコーダ-デコーダ型Transformerモデルである。このモデルは、専門家による反復学習を通じて微調整され、A^*探索よりも少ない探索ステップで最適なプランを生成する。我々の訓練方法では、A^*の探索ダイナミクスは、記号的プランニング中にタスク状態が探索木に追加・削除されるタイミングを概説するトークンシーケンスとして表現される。迷路ナビゲーションに関するアブレーション研究では、Searchformerが最適なプランを直接予測するベースラインを大幅に上回り、モデルサイズが5～10分の1、訓練データセットが10分の1で優れた性能を発揮することがわかった。また、Searchformerがソコバンのようなより大規模で複雑な意思決定タスクにスケールアップし、解決率の向上と探索ダイナミクスの短縮を実現する方法も示す。

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.