Title: Learning from Imperfection: Mistake-Aware LLM Finetuning for Robust Planning
Primary area: applications to robotics, autonomy, planning
Keywords: Large Language Models, Robust Planning, Embodied Agents, Mistake-Aware Finetuning

Abstract: While finetuned Large Language Models (LLMs) for embodied planning excel at producing reliable plans for the given environment, they do so in a very narrow area of operation. Usually one wrong step in such a plan is sufficient to get the agent into an unseen scenario. Current training paradigms focus on preventing mistakes by learning from optimal demonstrations, but neglect the crucial skill of recovering when deviating from the correct plan. To address this gap, we propose Mistake-Aware Finetuning (MAF), a novel training methodology that explicitly teaches agents to recover from planning errors. In the MAF paradigm, the model is exposed to plans containing intentional mistakes, but a targeted loss mask ensures it only learns from the subsequent, correct recovery actions. This allows the model to learn the association between a failure state and its resolution without being negatively influenced by the erroneous action itself.
We demonstrate the effectiveness of MAF by finetuning a Llama-3B model on two distinct environments. Our approach substantially improves the task success rate from 21% → 72% on the Tower of Hanoi puzzle and from 65% → 94% on the complex MiniGrid MiniBossLevel. Furthermore, to probe the generalization capabilities of our method, we introduce Unlock-To-Unlock-N, a novel and challenging benchmark designed to test long-horizon planning. We demonstrate that our MAF-trained agent not only performs robustly in this new benchmark but also exhibits generalization capabilities that extend beyond the training dataset. It highlights the MAF's immense potential for developing lightweight, robust embodied agents.

TL;DR: Finetuning LLMs with mistake-containing inputs, while excluding them from the loss, teaches models to recover from errors effectively.