Can language model plan in extrapolated environments?: Casestudy in textualized Gridworld

While language models have demonstrated impressive capabilities across generalized language tasks, their ability to extrapolate in a certain task is highly unknown. We first introduce the optimal path planning task in a textualized Gridworld environment as a valid probe for estimating the extrapolability of language models. We show that the mere next token prediction inherently fails to extrapolate in solving the task. Inspired by human cognition, we claim that language models should construct an internal simulation that explores the environment, i.e. cognitive map before actually interacting with the given environment. We demonstrate that auto-regressive generation of cognitive map and planning sequence can significantly enhance the performance of the planning power even in extrapolated environments, suggesting the necessity of cognitive map for language models as a path forward.