A Close Look At World Model Recovery In Supervised Fine-Tuned LLM Planners

Technical Summary: A Close Look At World Model Recovery in Supervised Fine-Tuned LLM Planners

Problem Statement
While Supervised Fine-Tuning (SFT) has improved the end-to-end planning capabilities of Large Language Models (LLMs), it remains unclear whether these models genuinely learn to represent and reason about the planning problems they solve. Specifically, do fine-tuned LLMs recover a "world model"—an internal understanding of state transitions and action validity—or do they merely learn to mimic valid action sequences? Previous research suggests LLMs struggle with planning in zero-shot settings and may learn shortcut solutions rather than robust reasoning mechanisms. This work investigates whether SFT on valid plan examples induces world model recovery in classical planning domains.

Methodology
The authors introduce a framework to evaluate world model recovery using two complementary definitions:

1. Internal Representations: Whether the LLM learns linear representations that encode the truth values of state predicates and the validity of actions.
2. Generative Capabilities: Whether the LLM assigns high probability to valid action sequences and low probability to invalid ones.

The study utilizes the gemma2-9b-instruct model, fine-tuned to generate end-to-end plans (dubbed "Plansformers") in two classical planning environments: Blocksworld and Logistics. The authors experiment with three distinct training data distributions to determine their impact on recovery:

• Optimal Plans: Generated by a heuristic solver (Fast Downward).
• Random Walk Plans: Generated by sampling random valid actions until a maximum step count is reached.
• Random Walk++ Plans: A modified random walk that invalidates actions leading to previously visited states to prevent cycles and improve state space coverage.

Additionally, the authors explore State Chain-of-Thought (State-CoT), where intermediate states are interleaved with actions in the training sequence to enhance groundedness.

To measure recovery, the authors employ:

• Linear Probes: Logistic regression models trained on hidden layer activations to predict action validity and state predicate truth values.
• Generative Classification: A metric evaluating whether the model's token probabilities rank valid actions higher than invalid actions when prompted with partial plans.
• Out-of-Distribution (OOD) Testing: Evaluation on planning instances with increased complexity (e.g., more blocks) than seen during training.

Key Results

1. Linear Encoding of Validity: Supervised fine-tuning on valid action sequences enables LLMs to linearly encode action validity in their internal representations. Probe F1 scores reach near-perfect levels (~1.0) in intermediate and upper layers across all training distributions.
2. Partial State Representation: Models learn to linearly encode the truth values of some state predicates (e.g., arm-empty, at ?pkg ?loc), but performance varies significantly by predicate complexity. Simpler predicates with fewer arguments are encoded more accurately than complex ones.
3. Divergence Between Internal and Generative Metrics: A critical finding is the mismatch between internal representations and generative outputs. Models fine-tuned on optimal plans possess high-quality internal representations of action validity (high probe scores) but fail to use token probabilities to distinguish valid from invalid actions during generation. Conversely, models fine-tuned on random walk data show superior generative classification (ranking valid actions higher than invalid ones) but struggle with OOD generalization.
4. Impact of State Space Coverage: Random Walk++ data, which ensures broader state space coverage by eliminating cycles, yields the most robust world model recovery. It achieves superior performance in both in-distribution generative classification and OOD planning tasks compared to standard random walks or optimal plans.
5. OOD Generalization: While internal representations of action validity remain robust to distribution shifts (OOD data), generative performance collapses for models trained on standard random walks. Models trained on Random Walk++ and Optimal plans show better resilience to OOD shifts, though performance still degrades as problem complexity increases.

Significance and Claims
The paper claims to provide a "recipe" for applying interpretability techniques to planning LLMs, offering a holistic view that combines internal representation analysis with generative evaluation. The authors conclude that:

• SFT on valid plans is sufficient for LLMs to learn internal representations that linearly encode action validity and certain state predicates, suggesting a form of world model recovery.
• However, the presence of internal representations does not guarantee that the model will utilize them effectively in its generative probabilities; the training distribution (specifically state space coverage) heavily influences this alignment.
• The findings are "optimistic" regarding the viability of fine-tuning to enhance planning, demonstrating that LLMs can partially recover planning world models. However, the work also highlights limitations, noting that the causal link between these representations and the model's planning process remains unproven, and that generalization to significantly larger parameter counts or natural language planning remains an open question.

The authors explicitly state they do not claim to have solved OOD planning or proven causal reasoning, but rather provide empirical evidence that fine-tuned LLMs develop structured internal representations of planning dynamics that are distinct from, yet related to, their generative capabilities.