Más allá de A*: Mejor planificación con Transformers mediante dinámicas de búsqueda Arranque inicial (Bootstrapping)

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

February 21, 2024

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

cs.AI

Resumen

Si bien los Transformers han permitido avances significativos en diversos entornos de aplicación, estas arquitecturas aún se quedan atrás de los planificadores simbólicos tradicionales para resolver tareas complejas de toma de decisiones. En este trabajo, demostramos cómo entrenar Transformers para resolver tareas de planificación complejas y presentamos Searchformer, un modelo Transformer que resuelve óptimamente rompecabezas de Sokoban previamente desconocidos en un 93.7% de los casos, utilizando hasta un 26.8% menos de pasos de búsqueda que el algoritmo A^* estándar. Searchformer es un modelo Transformer de codificador-decodificador entrenado para predecir la dinámica de búsqueda de A^*. Este modelo se ajusta mediante iteraciones expertas para realizar menos pasos de búsqueda que A^* mientras sigue generando un plan óptimo. En nuestro método de entrenamiento, la dinámica de búsqueda de A^* se expresa como una secuencia de tokens que describe cuándo los estados de la tarea se agregan y eliminan del árbol de búsqueda durante la planificación simbólica. En nuestros estudios de ablación sobre navegación en laberintos, encontramos que Searchformer supera significativamente a los modelos de referencia que predicen el plan óptimo directamente, con un tamaño de modelo 5-10 veces menor y un conjunto de entrenamiento 10 veces más pequeño. También demostramos cómo Searchformer escala a tareas de toma de decisiones más grandes y complejas, como Sokoban, con un porcentaje mejorado de tareas resueltas y una dinámica de búsqueda acortada.

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.