Source-linked AI summary
Benefits and Pitfalls of Reinforcement Learning for Language Model Planning: A Theoretical Perspective
Siwei Wang, Yifei Shen, Haoran Sun, Shi Feng, Shang-Hua Teng, Li Dong, Yaru Hao, Wei Chen
TL;DR
The theoretical basis for RL’s advantage over SFT in language-model planning remains unclear. The paper analyzes learning dynamics using a graph-based planning abstraction and examines PG and Q-learning. It identifies SFT’s co-occurrence solutions and PG’s diversity collapse, while Q-learning offers diversity preservation and off-policy learning under the studied conditions.
Problem
The theoretical basis for RL’s advantage over SFT in language-model planning, and the limitations of current RL methods, remain to be established.
Source: §1 · p. 1Method
The paper analyzes PG and Q-learning learning dynamics in a graph-based path-planning abstraction.
Source: §2.1 · p. 2, §1 · p. 2Results
Takeaways & Limitations
The findings clarify mechanisms behind RL-based planning and identify Q-learning as an approach that can combine diversity with off-policy learning once reward design addresses Q-value bias.
Source: §1 · p. 2, §6 · p. 9The theoretical analysis assumes persistent exploration and that next-token logits depend on the target and current nodes.
Source: §3.2 · p. 4, §5.1 · p. 8
Abstract
from arXiv · showhide
Recent reinforcement learning (RL) methods have substantially enhanced the planning capabilities of Large Language Models (LLMs), yet the theoretical basis for their effectiveness remains elusive. In this work, we investigate RL's benefits and limitations through a tractable graph-based abstraction, focusing on policy gradient (PG) and Q-learning methods. Our theoretical analyses reveal that supervised fine-tuning (SFT) may introduce co-occurrence-based spurious solutions, whereas RL achieves correct planning primarily through exploration, underscoring exploration's role in enabling better generalization. However, we also show that PG suffers from diversity collapse, where output diversity decreases during training and persists even after perfect accuracy is attained. By contrast, Q-learning provides two key advantages: off-policy learning and diversity preservation at convergence. We further demonstrate that careful reward design is necessary to prevent Q-value bias in Q-learning. Finally, applying our framework to the real-world planning benchmark Blocksworld, we confirm that these behaviors manifest in practice.
1 INTRODUCTION
The paper addresses the unclear theoretical basis for RL’s advantage over SFT in language-model planning using a tractable graph-planning analysis. It characterizes SFT’s spurious co-occurrence solutions, PG’s diversity collapse, and Q-learning’s diversity and off-policy advantages.
1 INTRODUCTION
- SFT can introduce co-occurrence-based spurious solutions, while RL methods achieve correct planning primarily through exploration.The analysis is intended to explain why SFT tends to memorize while RL promotes generalization.Source: §1 · p. 2
- PG’s exploration-driven data augmentation helps it outperform SFT, but its output diversity continues to decline even after training accuracy reaches 100%.KL regularization preserves diversity at the expense of accuracy.Source: §1 · p. 2
- With process rewards addressing Q-value bias, Q-learning preserves output diversity at optimal training accuracy and supports off-policy learning.The paper notes that quantized-model or large-batch rollouts are effectively off-policy.Source: §1 · p. 2
- The paper provides a theoretical treatment of RL’s effects on language-model planning and validates its findings through experiments.It applies the framework to Blocksworld, according to the paper’s abstract.Source: §1 · p. 2
2 PRELIMINARIES
The paper models planning as finding paths through a directed graph and uses this abstraction to define its SFT and RL data. It then specifies PG and Q-learning objectives and reward settings for the analysis.
2.1 PATH PLANNING DATASET: SYNTAX AND DATA SOURCES
- Planning is represented as path finding on an unknown directed graph, with nodes as tokens, edges as transitions, and reachability defined by the existence of a path.The Blocksworld example maps each configuration to a node and each valid move to an edge; the abstraction removes language semantics while retaining the problem difficulty.Source: §2.1 · pp. 2–3
- Reachable source-target pairs are split into train and test sets; SFT uses sampled valid paths, while RL generates sequences from pairs and receives outcome or step rewards.Test pairs are unseen during both SFT and RL training.Source: §2.1 · p. 3
2.2 REINFORCEMENT LEARNING ALGORITHMS
- PG uses an end-of-path outcome reward and a loss with KL regularization, while Q-learning approximates its Q-function with model logits over states and actions.The KL strength is controlled by λ; the study considers outcome rewards and process rewards, which penalize non-adjacent moves and reward reaching the target.Source: §2.2 · pp. 3–4
3 LIMITATIONS OF SUPERVISED FINE-TUNING IN PLANNING
The section characterizes SFT’s stable behavior and relates it to incomplete graph learning. Its analysis attributes the limitation to memorizing observed co-occurrences rather than exploiting transitivity, while RL approaches capture adjacency better in the reported comparison.
3 LIMITATIONS OF SUPERVISED FINE-TUNING IN PLANNING
- The stable-point analysis identifies SFT’s fundamental planning limitation and provides a theoretical basis for comparing it with RL’s exploration-driven generalization.The section’s characterization builds on earlier findings that SFT lacks mechanisms for learning transitivity.Source: §3 · p. 4
3.1 DISCUSSIONS ON EXISTING FINDINGS
- Prior analysis found that SFT may encode incomplete adjacency and reachability structures, and the Blocksworld comparison shows that even observed low-frequency edges may be poorly captured.Figure 1 compares training-edge frequencies with learned adjacency weights; the cited passage reports that some observed relations, especially low-frequency ones, are not well captured.Source: §3.1 · p. 4, Figure 1 · p. 5
3.2 CHARACTERIZATION OF THE STABLE POINT IN SFT-BASED LEARNING DYNAMICS
- Under the stated expressiveness assumption, the SFT optimum is characterized by training co-occurrence counts for target, current, and next nodes.When a target-current combination has no observed next-node counts, the output can be any probability distribution.Source: §3.2 · pp. 4–5
- SFT therefore memorizes dataset co-occurrences and fails to use transitivity information absent from the training data to capture true graph connectivity.The section contrasts this characterization with RL methods, whose learned adjacency structures capture relationships better.Source: §3.2 · p. 5
4 PATH PLANNING CAPACITIES OF POLICY GRADIENT
Policy gradient (PG) benefits from exploration-driven data expansion but can lose output diversity, while KL regularization preserves diversity at a potential cost to accuracy. Empirical results show that this trade-off also affects test accuracy during training.
4.1 THEORETICAL ANALYSIS
- With 0-1 outcome rewards and no KL regularization, each PG update is equivalent to SFT on correct paths generated during that iteration.Unlike fixed-data SFT, PG generates on-policy data and can discover new correct paths absent from the initial training set.Source: §4.1 · p. 6
- Without KL regularization, PG decreases logits for invalid next-node choices, driving the probability of wrong paths on training pairs toward zero.The analysis assumes Assumption 3.1 and 0-1 outcome rewards.Source: §4.1 · p. 6
- PG’s diversity collapse continues after it reaches 100% training accuracy, as output probabilities move away from a uniform distribution over valid next nodes.The theorem considers 0-1 outcome rewards, no KL regularization, and logits of invalid choices fixed at negative infinity.Source: §4.1 · p. 6
- KL regularization preserves some diversity by keeping the policy close to the base model, but can limit training accuracy and potential gains.Its benefit depends on the base model: it may improve generalization when that model is capable, but hinder learning when its prior is weak.Source: §4.1 · p. 7
4.2 EMPIRICAL VALIDATIONS
- Without KL regularization, PG loses output diversity during training, and test accuracy degrades as diversity diminishes.PG can maintain 100% training accuracy while diversity keeps falling until only one path is produced per source–target pair.Source: §4.2 · p. 7, Figure 2 · p. 7
5 ANALYSIS OF THE Q-LEARNING-BASED PLANNING MECHANISM
Q-learning’s outcomes depend on reward design: process rewards preserve graph structure and diversity, whereas outcome-only rewards can collapse to biased solutions. The analysis and experiments also support off-policy learning with process rewards.
5.1 THEORETICAL ANALYSIS
- With outcome rewards, Q-learning’s stable logits are constant across states for each target, losing structural information about the graph.Under persistent exploration and Assumption 3.1, the result applies to next nodes other than the target.Source: §5.1 · pp. 8–9
- Process rewards encode adjacency and target conditions, allowing Q-learning logits to converge to values reflecting the graph’s structure.The convergence analysis assumes Assumption 3.1, zero-initialized scores, and persistent exploration; the rate depends on the learning rate and update proportions.Source: §5.1 · pp. 8–9
- Process rewards mitigate the Q-value bias that can arise when Q-learning relies only on outcome rewards.The analysis also extends process-reward results to a simplified linear Transformer without Assumption 3.1.Source: §5.1 · pp. 8–9, §5 · p. 8
- Under persistent exploration, process-reward Q-learning converges while preserving action diversity and remains valid under off-policy sampling.The persistent-exploration condition requires every coordinate to be updated frequently enough; epsilon exploration satisfies it.Source: §5.1 · pp. 8–9
5.2 EMPIRICAL VALIDATIONS
- Process-reward Q-learning achieves comparable training accuracy and significantly better test accuracy than PG, while outcome-reward Q-learning converges to near-zero accuracy on both sets.Off-policy process-reward Q-learning attains training and test accuracy comparable to its on-policy counterpart, and process rewards preserve output diversity.Source: §5.2 · p. 9
6 CONCLUSION
The conclusion presents a theoretical account of how exploration, reward design, and learning method shape language-model planning. It identifies path planning over graphs as a useful case study and points to Q-learning as a direction for future work.
6 CONCLUSION
- The paper’s analysis attributes SFT’s spurious co-occurrence solutions, RL’s advantage, and PG’s diversity collapse to learning dynamics and exploration.It identifies Q-learning as mitigating diversity collapse while supporting off-policy learning.Source: §6 · p. 9
- The authors identify Q-learning as a direction for robust, scalable, and generalizable language-model planning.They point to combining diversity and accuracy and enabling off-policy learning as future directions.Source: §6 · pp. 9–10
B.1 PLANNING OF LLMS
- The paper frames planning as graph path-finding, with examples including tool dependencies, theorem dependencies, and sequences of game skills.The examples motivate the paper’s path-planning abstraction.Source: §B.1 · p. 14
B.2 RL FOR LLMS
- Prior accounts of RL’s benefits include empirical generalization comparisons, theoretical analysis of verification-free methods, and an entropy–accuracy trade-off.The paper positions its path-planning case study alongside these existing accounts.Source: §B.2 · p. 14
B.3 GRAPH PROBLEMS WITH LANGUAGE MODELS
- The paper describes gradient-dynamics analysis as a practical but challenging approach, contrasting it with empirical interpretability and transformer expressiveness analyses.It characterizes its work as, to the authors’ knowledge, the first analysis of RL gradient dynamics in LLMs.Source: §B.3 · p. 15
C.1 PATH PLANNING ALGORITHM IN TRANSFORMER
- The handcrafted path-planning algorithm samples a next node from those adjacent to the current node and able to reach the target.It appends sampled nodes to the path until reaching the target, then outputs the path.Source: §C.1 · p. 15
C.2 PROOF OF THEOREM 3.1
The proof rewrites next-token prediction loss using counts of observed target, current, and next-node tuples. It identifies the resulting cross-entropy minimizer and notes that unobserved pairs impose no loss constraint.
C.2 PROOF OF THEOREM 3.1
- Under the assumption that the output distribution depends only on target and current nodes, identical terms in the loss can be aggregated.The aggregation is expressed using counts of observed tuples.Source: §C.2 · p. 15
- When tuples are observed, the bracketed loss is cross-entropy between the empirical distribution and the softmax of the model’s score vector.The loss is minimized when these distributions match.Source: §C.2 · pp. 15–16
- When no next-node tuple is observed for a target–current pair, the loss does not depend on that pair’s output distribution.The distribution can therefore be any valid probability distribution.Source: §C.2 · p. 16
D.1 PROOF OF THEOREM 4.1
The proof excerpts derive gradient and stationarity conditions for the learning updates, then establish persistent exploration and a shared stable-point value for non-target outputs.
D.2 PROOF OF THEOREM 4.2
- The derivations use sampled target–current–next-node counts to form gradients and characterize stable points through vanishing expected updates.The excerpts include both on-policy count expectations and stable-point conditions.Source: §D.2 · p. 16, §D.4 · pp. 17–18, §E.2 · p. 18
E.1 PROOF OF LEMMA 5.1
- Uniform ε-exploration gives every target–current–next-node triple positive probability, ensuring each corresponding score coordinate is updated.Each occurrence of a triple triggers an update, yielding the persistent exploration condition.Source: §E.1 · p. 18
E.2 PROOF OF THEOREM 5.1
- At stable points, for fixed target i and k ≠ i, all f(i,j)[k] share a value independent of j and k.The proof obtains this common value from the stationarity conditions.Source: §E.2 · pp. 18–19
E.3 PROOF OF THEOREM 5.2
The process-reward recursion converges to casewise limits tied to graph structure, with convergence rates shaped by the path of dependencies and update frequency.
E.3 PROOF OF THEOREM 5.2
- The recursion’s casewise limits depend on adjacency A[j,k] and whether k is an ancestor of target i.The displayed cases converge to the theorem’s stated limits.Source: §E.3 · pp. 19–20
- Along a directed path, each node’s error can decay only after the next node’s error has decayed, yielding a product of per-edge contraction rates.The proof describes the effective contraction factor as the product of the per-edge rates.Source: §E.3 · p. 20
- All iterates converge linearly in global time, with effective rates jointly determined by η and the update proportions.Under persistent exploration, the relevant coordinate is updated with positive frequency.Source: §E.3 · p. 20
E.4 PROOF OF THEOREM 5.3
The proof characterizes Q-learning stable points using limiting sampled-transition frequencies and graph reachability; it also establishes the gradient equivalence of unclipped PPO and vanilla policy gradient.
E.4 PROOF OF THEOREM 5.3
- Positive limiting frequencies yield stationary conditions whose solutions encode the graph’s reachability relation, up to a one-dimensional shift for each k.The reachability recursion has the unique {0,1} solution R[i,k], while the solution set is parameterized by c_k.Source: §E.4 · pp. 20–21
- The characterized stable-point forms are sufficient as well as necessary: substituting them into the update equations makes the expected increment vanish.The proof states that the expected gradient equations cancel by construction.Source: §E.4 · p. 21
F EQUIVALENCE OF UNCLIPPED PPO AND POLICY GRADIENT
- For any fixed sequence, unclipped PPO with a stop-gradient denominator has the same gradient as vanilla policy gradient.The denominator is treated as constant, so its gradient vanishes.Source: F EQUIVALENCE OF UNCLIPPED PPO AND… · p. 21
G.1 VALIDATION OF LEARNED ATTENTION
Attention visualizations examine learned attention during training and show target-node focus alongside access to current-node information in one- and two-layer transformers.
G.1 VALIDATION OF LEARNED ATTENTION
- In SFT, target-node attention peaks quickly and then declines while remaining dominant; the authors hypothesize this may reflect overfitting and auxiliary prediction strategies.They connect this possible overfitting to SFT’s decreasing generalization performance in their experiments.Source: §G.1 · p. 22
- In PG and Q-learning, target-node attention increases throughout training, exceeding 95% at the end for Q-learning.The final Q-learning attention is closest to the conditions in Assumption 3.1.Source: §G.1 · p. 22
- In the two-layer, one-head transformer, both layers predominantly attend to target and current nodes, consistent with the same behavior across architectures.The attention map is averaged over the SFT dataset.Source: §G.1 · p. 22, Figure 6 · p. 23
G.2 EXPERIMENTS ON ERD ˝OS-R´ENYI GRAPHS
The experiments extend SFT-initialized PG and Q-learning to split-based Erdős–Rényi graph tests, then assess learned adjacency on graph-modeled four-block Blocksworld. They find split-dependent effects of KL regularization, slower Q-learning convergence on new training pairs, and improved adjacency under both RL methods in Blocksworld.
G.2 EXPERIMENTS ON ERD ˝OS-R´ENYI GRAPHS
- After SFT, reachable pairs are split into RL train and test sets, yielding intersections that measure learning on new pairs and forgetting on withheld SFT pairs.Source: §G.2 · p. 23
- λ = 10−4 achieves the best balance: weaker KL regularization improves learning on new pairs, while no KL tends to overfit and forget SFT knowledge.Stronger KL regularization hinders learning new pairs, and overly strong regularization can also hinder PG’s improvement on DTest2Test.Source: §G.2 · p. 24
- Q-learning converges more slowly than PG in this setting; the authors suggest poor initial performance on new pairs and more failure cases may cause stronger re-instantiation.Source: §G.2 · p. 24
G.3 EXPERIMENTS ON GRAPH REPRESENTED FOR BLOCKSWORLD
- For four-block Blocksworld, valid configurations form a 73-node undirected graph, and SFT uses 50,000 sampled paths with source and target nodes drawn uniformly.Source: §G.3 · p. 24
- On Blocksworld, both PG and Q-learning improve learned graph adjacency over SFT, with Q-learning nearly recovering the complete adjacency.The metric measures the model’s understanding of graph adjacency.Source: §G.3 · p. 25
Related papers
This paper cites
- Voyager: An Open-Ended Embodied Agent with Large Language Models2023
- HybridFlow: A Flexible and Efficient RLHF Framework2024
- Playing Atari with Deep Reinforcement Learning2013
- Progress measures for grokking via mechanistic interpretability2023
- PlanBench: An Extensible Benchmark for Evaluating Large Language Models on Planning and Reasoning about Change2022