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

TL;DR

Spatial reasoning in LLMs is challenged by multi-step planning in structured spaces. The authors propose a two-stage framework: first supervised fine-tuning on atomic spatial transformations to build physics priors, then freezing the backbone and applying GRPO-based reinforcement learning over LoRA adapters to compose primitives into plans. Evaluations on an ASCII-art spatial task in Dynamic and Static settings show that the proposed method outperforms baselines and converges faster, with interpretable changes in attention patterns. This modular, reusable building-block approach advances spatial understanding and planning in LLMs and offers a scalable path toward more interpretable reasoning.

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.

Figures (5)

• Figure 1: Overview of the proposed training framework. The pipeline follows a two-stage learning approach. In the SFT Phase, the base model (Qwen-Instruct) is fine-tuned on the Building Block Dataset to acquire atomic spatial priors (translation, scaling, rotation), resulting in the intermediate Qwen-Physics model. In the RL Phase, we employ GRPO in a closed-loop setting with LoRA adapters. The model is trained to master multi-step spatial reasoning and planning, yielding the final Qwen-PhysRL model.
• Figure 2: GRPO training reward trajectories for Qwen-PhysRL and Qwen-DirectRL, illustrating improved stability and faster convergence when using a frozen physics-aware prior.
• Figure 3: Left: Step-by-step cumulative reward comparison between Qwen-PhysRL (Dynamic and Static settings) and random policy ($\pi_{rnd}$) on tasks requiring 3 translations, 1 rotation, and 1 scaling. Right: Action-reward trajectory illustrating the optimal path (blue: scale$\rightarrow$tr$\rightarrow$rot$\rightarrow$tr$\rightarrow$tr), which achieves the maximum reward of 2.7, contrasted with alternative action sequences (gray). The trajectory obtained under the Static setting is shown in red, where the model initially follows the optimal prefix (scale$\rightarrow$tr$\rightarrow$tr$\rightarrow$tr) but then diverges and loses track of the remaining steps, failing to complete the full plan.
• Figure 4: Token-level attention distribution across layers for Qwen-Physics and Qwen-Instruct.
• Figure 5: Layer-wise attention score analysis showing (a) attention divergence of Qwen-Physics from the backbone Qwen-Instruct model across layers, and (b) relative attention allocation between system prompt and map tokens.