作者： Lucas Lehnert, Sainbayar Sukhbaatar, DiJia Su, Qinqing Zheng, Paul Mcvay, Michael Rabbat, Yuandong Tian
提交/修订日期： 2024年2月21日提交，2024年4月26日修订
主题/分类： 人工智能 (cs.AI)
摘要：
尽管Transformer架构在各种应用场景中取得了巨大进展，但在解决复杂决策任务方面，此类架构仍然落后于传统的符号规划器。在这项工作中，我们展示了如何训练Transformer来解决复杂的规划任务。这是通过训练一个编码器-解码器Transformer模型来预测A搜索算法的搜索动态实现的。我们微调该模型，得到一个名为Searchformer的Transformer模型，该模型能以93.7%的成功率最优地解决之前未见过的Sokoban谜题，同时使用的搜索步骤比最初用于训练的A实现最多减少26.8%。在我们的训练方法中，A*的搜索动态被表达为一个标记序列，该序列概述了在符号规划过程中任务状态何时被添加和移除到搜索树中。Searchformer显著优于直接预测最优规划的基线模型，其模型规模小5-10倍，训练数据集小10倍。最后，我们展示了Searchformer如何扩展到更大、更复杂的决策任务，并提高了任务解决率和缩短了搜索动态。