За пределами A*: Улучшенное планирование с помощью трансформеров через динамику поиска Бутстреппинг

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, которая оптимально решает ранее не встречавшиеся головоломки Sokoban в 93,7% случаев, при этом используя до 26,8% меньше шагов поиска по сравнению со стандартным алгоритмом A^*. Searchformer представляет собой модель Transformer с архитектурой "кодировщик-декодировщик", обученную предсказывать динамику поиска алгоритма A^*. Затем эта модель дообучается с использованием экспертных итераций, чтобы выполнять меньше шагов поиска, чем A^*, при этом всё ещё генерируя оптимальный план. В нашем методе обучения динамика поиска A^* выражается в виде последовательности токенов, описывающей, когда состояния задачи добавляются и удаляются из дерева поиска в процессе символического планирования. В наших исследованиях по навигации в лабиринтах мы обнаружили, что Searchformer значительно превосходит базовые модели, которые предсказывают оптимальный план напрямую, при этом размер модели в 5-10 раз меньше, а объём обучающих данных — в 10 раз меньше. Мы также показываем, как Searchformer масштабируется на более крупные и сложные задачи принятия решений, такие как Sokoban, с улучшенным процентом решённых задач и сокращённой динамикой поиска.

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.