From Building Blocks to Planning: Multi-Step Spatial Reasoning in LLMs with Reinforcement Learning

TL;DR

This paper introduces a two-stage method combining supervised fine-tuning and reinforcement learning to improve multi-step spatial reasoning in large language models, demonstrating superior performance and stability in puzzle environments.

Contribution

The paper presents a novel two-stage approach that decomposes spatial reasoning into atomic blocks and their composition, enhancing LLMs' multi-step planning abilities.

Findings

Outperforms baseline models in spatial reasoning tasks

Faster convergence and more stable training than end-to-end RL

Attention analysis shows improved spatial understanding

Abstract

Spatial reasoning in large language models (LLMs) has gained increasing attention due to applications in navigation and planning. Despite strong general language capabilities, LLMs still struggle with spatial transformations and multi-step planning in structured environments. We propose a two-stage approach that decomposes spatial reasoning into atomic building blocks and their composition. First, we apply supervised fine-tuning on elementary spatial transformations, such as rotation, translation, and scaling, to equip the model with basic spatial physics. We then freeze this physics-aware model and train lightweight LoRA adapters within the GRPO framework to learn policies that compose these building blocks for multi-step planning in puzzle-based environments, in a closed-loop manner. To support this pipeline, we synthesize an ASCII-art dataset and construct a corresponding ASCII-based reinforcement learning environment. Our method consistently outperforms baselines, including the generic backbone, physics-aware model, and end-to-end RL models, under both Dynamic environments with explicit state updates and Static environments where the model must rely on its internal state across steps. In addition, the proposed approach converges faster and exhibits more stable training compared to end-to-end reinforcement learning from scratch. Finally, we analyze attention patterns to assess whether fine-tuning induces meaningful improvements in spatial understanding.