Benefits and Pitfalls of Reinforcement Learning for Language Model Planning: A Theoretical Perspective
Authors: Siwei Wang
Yifei Shen
Haoran Sun
Shi Feng
Shang-Hua Teng
Li Dong
Yaru Hao
Wei Chen
Venue: ICLR 2026 Poster (2026)
Keywords: LLM Planning
Path Planning
Reinforcement Learning
Primary area: interpretability and explainable AI

Abstract: 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.

Published as a conference paper at ICLR 2026
BENEFITS AND PITFALLS OF REINFORCEMENT LEARNING FOR LANGUAGE MODEL PLANNING: A THEORETICAL PERSPECTIVE
Siwei Wang"1, Yifei Shen'l, Haoran Sun Shi Feng Shang-Hua Teng", Li Dong', Yaru Hao', Wei Chen's
Microsoft Research Asia, "Peking University, "Harvrrd University, 4 Uniiversity of Southern California
ABSTRACT
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 bene- fits and limitations through a tractable graph-based abstraction, focusing on pol- icy gradient (PG) and Q-learning methods. Our theoretical analyses reveal that supervised fine-tuning (SFT) may introduce co-occurrence-baeed spurious solu- tions, whereas RL achieves correct planning primarily through exploration, un- derscoring exploration's ' role in enabling better generalization. However, 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 off-policy learning and diversity preser- vation at convergence. We further demonstrate that careful reward design is nec- essary 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.
INTRODUCTION
Planning is fundamental cognitive construct that underpins human intelligence, shaping our abil- ity to organize tasks, coordinate activities, and formulate complex solutions such as mathematical proofs. It enables humans to decompose complex goals into manageable steps, anticipate potential challenges, and maintain coherence during problem solving. Similarly, planning plays pivotal role in state-of-the-art Large Language Models (LLMs), enhancing their ability to address structured and long-horizon tasks with greater accuracy and reliability.
Early generations of LLMs primarily relied on next-token prediction and passive statistical learning, which limited their planning capabilities to short-horizon, reactive responses. The ol family of mod- els represents a major advance in planning by incorporating reinforcement learning (RL) objectives that reward accurate, multi-step reasoning and penalize errors. Inspired by the success of ol, RL has been applied to enhance planning capabilities in various settings, including task decomposition for tool (Wu al., 2024a Luo et al. 2025) and gaming (Yang al., 2024). visual-language spatial navigation (Chu et al.) 2025), and long horizhnoozon robotics tasks (Dalal et al. 2024). These ap- proaches have demonstrated significantly better performance than their supervised fine-tuning (SFT) counterparts. For more related works, please refer to Appendix B
Despite recent successes, the thcoretical basis underlying RL's advantage over SFT in planning tasks and the limitations of current RL methods remain to be established. To enable tractable analysis of the gradient dynamies, we adopt the data generation model from (Wang et al. (024b). Within their framework, planning is abstracted tool-use scenario can be modeled as identifying valid call sequence within an API call graph (Wu
path-finding problem over graph structure. For example,
et al. 2024b).
*t denotes equal contribution. Corresponding author (weic@@icrosoft.com).

---

Published as a conference paper at ICLR 2026
To capture the fundamental limitations of SFT in planning. we begin by presenting structural characterization of its stable point for path planning (Section(3). Our analyses, expanding the obser- vation of Wang et al.] (2024b) that transformer-bas -based LLM architectures cannot identify reachability relationships through transitivity in SFT, show that it introduces co-occurrence-bant spurious solu- tions into planning tasks. This characterization provides basis for comparison with and motivation for using the RL-based learning approach in language model planning.
Focusing on the behaviors of RL-based learning dynamics, we first consider policy gradient (PG), a widely adopted algorithm for tuning large language models (Section 4). Our analysis yields three key findings. First, with only 0-1 outcome rewards, each iteration of PG equivalently corresponds to an SFT process on the exploration data; however, PG empirically outperforms SFT due to the exploration-driven data augmentation it enables. Second, although PG converges to a model that outputs correct paths for all source- pairs seen during training, we uncover diversity collapse phenomenon: the model's output diversity steadily declines throughout training and continues to diminish even after achieving 100% training accuracy. Third, we show that KL regularization acts as an explicit
term, but at the expense of accuracy.
We then analyze Q-learning, paradigm well known in game playing but rarely applied to LLMs (Mnih et al. 2013) (Section 55.. Our analysis yields two key findings. First, when trained with only outcome reward signal, Q-learning suffers from Q-value bias; however, incorporat- ing process rewards eliminates this issue. Second, once this issue is addressed, Q-learning offers two theoretical advantages over PG: it converges to solution that preserves output diversity when achieving optimal training accuracy, and it naturally supports off-policy learning. The latter is par- ticularly important in practice, since rollouts performed with a quantized model or large batch sizes are effectively off-policy, as exemplified by the VeRL framework (Sheng et al. 2024). Finally, we validate all these theoretical findings through experiments.
To summarize, our main contribution is theoretical treatment of the impact of reinforcement learn- ing on language model planning. Our mathematical analysis of learning dynamics sheds light on phenomena observed in practice-for example, SFT tends to memorize while RL promotes gen- eralization; PG methods often suffer from diversity collapse; and KL. regularization helps mitigate diversity degradation, albeit at the cost of reduced accuracy. Other findings point to promising future directions, such as leveraging Q-learning to achieve both diversity and accuracy, as well as enabling off- policy learning. Taken together, these results provide principled foundation for understanding and advancing reinforcement learning methods in language model planning.
2 PRELIMINARIES
2.1 PATH PLANNING DATASET: SYNTAX AND DATA SOURCES
Following (Wang et al. 2024b), we abstract planning in large language models path planning over an unknown directed graph G = (V,E), where V represents the set of nodes and & represents the set of edges. Each node o E V is represented by unique token. The language model's vocabulary consists of these node tokens and a special end-of-sequence token, \n. An edge (u, v) E E signifies node A node L is reachable from node if directed
directed connection from node
path from 8 to I exists in G. We denote by A E {0, 1)]Vi(|V| the adjacency matrix of G, where
A[u, v] = 1 if and only if (u, v) E E, and by R E {0, V|x[V| the reachability matrix, where R(t, s] if and only if t is reachable from s.
Running Example (Blocksworld). To connect this abstraction to real-world LLM planning sce- narios, consider the Blocksworld domain (Valmeekam al. 2023b). In Blocksworld, we are given several colored blocks (e.g., Grey, Black, Red, White) placed either on a table or stacked on each other, and the task is to transform an initial arrangement (source state) into a target arrangement (target state) using valid moves. For example, the source state may place all blocks on the table, and the target state requires stacking them so that Red is on Grey, Grey is on Black, and Black is on White. We map every distinct block configuration to node in V: edges in E correspond to valid single moves such as "place White on Grey". ". valid plan is therefore equivalent to a path in g connecting the source node and target node l.
2

---

Published as a conference paper at ICLR 2026
This abstraction matches natural language planning tasks: in the original benchmark, the LLM is given two textual descriptions of initial and target states and asked to generate a sequence of natural language actions achieve the goal. In our abstract setup, strip away language semantics to focus on the core planning structure while retaining the same problem difficulty.
The set of all reachable source target pairs (s, t) is partitioned into training Dirain and test set Dest We define three corresponding data stages:
We construct training dataset DSFT for supervised fine-tuning by sampling multiple (K) paths for each reachable pair (s, t) E Drrain by random walk. Each training data in DSPT is sequence in the format "s a e \n", where b e tokens for nodes in valid path from to i, and n indicates the end the sequence. We call the model after SFT training the base model.
SFT Training Data:
RL Training Data: We sample pairs (s, from Draain and let the model itself (on-policy) or the base model (off-policy) generate the remaining tokens in the sequence. When the model outputs in or the generation reaches the maximum length, an outcome reward or some step rewards will be given, depending the used reward format.
Test Data: When testing, we provide pairs (s,t) from Dreat which are never encountered in either SFT or RL training. The model is tasked with generating a valid path from s to t.
Throughout the empirical study, we a one-layer, single-head Transformer as the backbone model. The embedding size is set to d = 120. The graph G in our main empirical validation is generated using the Erdós-Rényi model with |V] 100 nodes and an edge probability of 0.15. The ratio of the sets |Train Drsst is approximately 0.25 (approximately 20% pairs are in Drrain). The number paths sampled for each reachable pair in Drrain K = 10. We also consider the graph GBW that characterizes the transition between different block configurations in Blocksworld, which is proposed by Valmeekam et al. (2023a) to evaluate the LLM's planning ability. The details for the graph construction are presented Appendix G
Dressl
2.2 REINFORCEMENT LEARNING ALGORITHMS
We first define the notation. Given a vector x, we denote its m-th element by x[m]]. For given sequence "s b t \n", we represent it im the output probability vector of the current model at the m-th position, and by á base that of the base model before RL. The model parameters are denoted by 0.
(Moouce Marget, " -um, ".m....). We denote by
Policy Gradient. Let P be the set of valid paths. The outcome reward is only given at the end of the path and is defined by R(u) duer + p. where O and p constants, and d denotes the indicator function that is if condition is true and 0 otherwise. For an individual trajectory, the loss function is
üm [2mm+]]
l -=
Σ m21
R(u) log üm [onuum+]] + log üm [um+1] 1] log
base (1mm 1
Policy Gradient
KI. Divergence
where À controls the KL regularization strength, and {.} means the term is detached and will not contribute the gradient.
Q-Learning. The goal of Q-learning is approximate the Q-function with the model logits. Let Q0 am) ) be the Q-function where Sm = (Msourco: Usaget: u um) is the state, Qm E V the action, and s'
(Usource: Mtagatt ". Um, am) is their next state. The objective is
Σm (Qo(sm am)) [R(sm (m am) + max Qs(s' dnn))) We denote the logits at step by Um. For an individual trajectory, the loss is given by
üm [um+1] - R(u, m)) - max üm 11111]
(2)
m2
3

---

Published as a conference paper at ICLR 2026
For Q-learning's reward R(u, m), we study two scenarios: (i) outcome reward, where the reward depends on whether the path is correct, and (ii) process reward, where intermediate rewards are given based on adjacency and target checks. Specifically,
due Oug Uuuger?
If outcome reward,
R(u, m)
0um+1 Mingg
(um tho+t )88.
If process reward.
Target check
Adjucency check
That is, in the outcome reward setting, a reward of 1 is given only the entire path is valid, and it assigned at the step when the target is reached. In contrast, in the process reward setting, we do not check whether the entire path is valid or not. The model is always rewarded upon reaching the target, but it is also penalized at any step that transitions non-adjacent node.
LIMITATIONS OF SUPERVISED FINE- TUNING IN PLANNING
Focusing on the stationarity of the training dynamics, we present a basic structural characterization that captures a fundamental limitation of SFT in planning. Our analysis builds on an early find- ing of [aan et al.) (2024b). which showed that transformer-baset SFT planning approaches lack transitivit--leaning mechanisms needed to obtain complete reachability structures. The new char- acterization expands and complements the earlier results and provides a theoretical explanation for why SFT-based planning tends to rely on memorization. More importantly, this result establishes a thcoretical basis for comparison with RL-based planning frameworks and highlights the role of exploration in achieving better generalization during the adaptive learning process.
3.1 DISCUSSIONS ON EXISTING FINDINGS
To set up our characterization, we first review the analysis framework of Wang et al. (2024b), which examines the training dynamics of a one-layer, single-head Transformer under an autoregressive loss function. Their analysis shows that, during training, the model encodes both the adjacency and reachability structures of the underlying graph in its learnable parameters. The model then predicts the next node in a sequence by ensuring that it is adjacent to the current node and lies along path toward the target node. A full description of their approach given in Algorithm T in Appendix C
Wang et al. (2024b) showed, both theoretically and experimentally, that the adjacency and reach- ability information stored in a model's weights is generally incomplete, To formalize this, con- sider training dataset DSFT The observed adjacency matrix (DSFFT) contains exactly those edges (j. k) that appear in at least one path from DSFT Similarly, the observed reachability matrix R (DSFT) records that target node t is reachable from intermediate node k if DSFT contains a sequence with target in which k occurs as a non-source node. We refer to such pairs (t, k) as observed reachable pairs.
However, we find that even when an adjacency relation appears in DSFT the SFT model may not leamn high weight for it. To illustrate this, we run experiments on the Blockworld dataset, and the results are presented in Figure [ In Figure Ia we show the frequency of all adjacency relation- ships in the training set (every adjacency relationship appears least once), where brighter regions indicate higher frequencies. Then Figure [b displays the corresponding weights learned after SFT. By comparing them, we observe that some adjacency relationships present the data are not well captured by the model, especially those with low frequency. This observation motivates us to further investigate the model's stable (optimal) points.
3.2 CHARACTERIZATION OF THE STABLE POINT IN SFT- SFT-BASED LEARNING DYNAMICS
Building on the observation of Wang al. 2024b) that next-node prediction depends mainly on the current and target nodes, we adopt the following natural assumption about model expressiveness for our structural characterization. Recall that Utarget and um denote the target node and the current node at position m, respectively.
Assumption 3.1. The model's predicted logits for the next token can be expressed as a function of the (target, current) node pair, i.e., there exists a function f such that the logits üm {(uargets Um
4

---

Published as a conference paper at ICLR 2026
61-
(a) Edge Frequency
(b) SFT
(c) SFT PG
(d) SFT + Q-Learning
Frequency of edge occurrences in the SFT training data DSFT and the adjacency structures learned by different models. The underlying graph represents transitions between block configura- tions in Blocksworld (Valmeekam et al. 2023a).
Figure 1:
Note that the assumption, f can be an arbitrary function. Our experiments validate this assump- tion. As shown in Section G. the evolution of attention maps during training for SFT, PG, and Q-learning demonstrates that the trained transformer acts primarily as function of the target and current nodes.
We now characterize the structure of the stable point achieved by SFT. Due to space limitations, we defer all the proofs in this paper the appendix.
Theorem 3.1 (Optimal Solution of SFT). Assume Assumption 3.1 holds. Let Nu.p denote the number of occurrences the truining dataset where the larget node Uhareet the current node is Um and the next node is k. The optimal solution of SFT satisfies:
exp( {( (Mharrger Vmm [kd]
Natages Mhrr de
if
N Munger Mom >> O.
Ene exp(f( (ehaarua Usarget, Um [kT)
N Nawager Maen Jc
EE N Mougee M K' 0, output can be any probability distribution.
Takeaway 1: SFT memorizes co-occurrence relationships in the training dataset.
Theorem 3.I extends the findings of Wang al. (024b), which showed that may fail to learn the complete adjacency and reachability matrices, leading to spurious correlations. However, those earlier results did not specify the nature of the solutions to which the model con- verges. Complementing their work, Theorem 3.I clarifies this by showing that SFT essentially memorizes Co- co-occurrence relationships among the target node, the current node, and the immediate next node based on their frequencies in DSPT Hence. SFT will fail to exploit transitivity information (which never appears in DSFT) to capture the true graph connectivity required for path planning.
mechanisms
further compare the weights of models trained by two RL approaches, PG and Q-learning. Both RL approaches capture the adjacency relationships better. Similar findings are re- ported by Chu al. (2025), who empirically observe that SFT tends memorize while RL exhibits better generalization. Our structural analysis in Theorem 3.1 provides a theoretical explanation for the first part of this phenomenon, namely, why "SFT memorizes". In the following sections, examine the two RL-based approaches, PG and Q-learning, and provide a theoretical explanation of the second part, i.e., why "RL generalizes".
In Figure [
4 PATH PLANNING CAPACITIES OF POLICY GRADIENT
In this section, we examine the path-planning capacity of the policy gradient, the core principle behind advanced RL algorithms such as PPO (Schulman et al. 2017) and GRPO (Shao et al. 2024). Understanding the strengths and limitations of the basic policy gradient provides theoretical insights into its behavior, highlights the mechanisms that enable effective path planning, and clarifies the challenges that motivate more sophisticated approaches.
5

---

Published as a conference paper at ICLR 2026
4.1 THEORETICAL ANALYSIS
We first establish the connection between policy gradient (PG) and supervised fine-tuning (SFT), highlighting the potential advantages of PG over SFT. We then analyze PG's training dynamics and show that, without KL regularization, the model can achieve 100% training accuracy (under temperature sampling) while progressively losing output diversity. Finally, we demonstrate that, when initialized with reasonably capable base model, adding a KL regularization helps preserve diversity and thereby enhances generalization, albeit sometimes the cost of accuracy.
To make this connection precise, we show that the PG loss function closely resembles the SFT loss, restricted to the subset data generated during RL training that corresponds to correct paths.
Theorem 4.1 (Connections between PG and SFT). Assume Assumption holds. Let DRL denote the set of data generated during the RL training step t. When = 1. p 0 (i.e., reward I for (ie., without KL. regularization), the loss function
correct puth and reward otherwise) and of Policy Gradient is the same as the loss function of using SFT only on correct paths in DRL,.
As shown by Wang et al. (2024b). SFT can learn the adjacency and reachability relations. Thus, Theorem 4.1 shows that PG can capture these relations presented in the dataset ( DRL..) n P. However, unlike SFT, which relies on fixed training dataset, PG generates data on on-policy during training. As the model improves, it can explore and discover new correct paths that were absent from the initial training set. This exploration-driven data augmentation enables PG to achieve stronger performance beyond what SFT alone can provide.
Takeaway 2: PG outperforms SFT primarily because its iterative data generation process encourages exploration and effectively expands the training dataset.
Building the loss function, analyze the gradient and identify two distinctive properties of on-policy PG updates.
Theorem 4.2 (Convergence of PG without KL regularization). Assume Assumption 3.hholss. For any i.J pair, let C(i, denote the set of nodes that can reach and are adjacent to The following B8 for k & C(i,j) is always then holds: If 1, O and X = O, then (i) the grudient (E(L,))R]
posilive, and (ii) the total sum of gradient Σx (F(
0t = O.
Theorem 4.2
shows that the logits [A]] corresponding to incorrect tuples (i.j, i.c., cases where node cannot reach node i through node k, will continue to decrease, while some other logits will not converge to -c.. Consequently, under gradient descent, the probability that the model outputs wrong path in D Dirain converges to zero.
Next, we analyze how the model's output diversity evolves. Intuitively, the most diverse model that still achieves 100% accuracy is that produces uniform distribution over C(i,j) for each target node i and current node j. We now analyze the evolution of the KI. divergence between this uniform distribution and the model's 's output distribution during PG training without KL regularization.
Theorem 4.3 (Diversity Collapse of PG without KL. regularization), Assume Assumption 3./ holds.
Let U denote the uniform probability distribution on support C(i,j). When = 1, = 0 and = O, and logits (kk for & (ii.j) is -cc, where (i,.)) denotes the logits value of f(i.j) time step I. For any such PG grudient descent step we have that
KL(Uo(,)) softmax (i, L[ |softmax (F ((,,))))
Note that the metric K LL(Uc(,,)
|| [|softmax(F' (i) takes minimum value when softmax((" also the uniform distribution on C(i,)), and takes maximum value when softmax(f" (f. is a one-hot vector. Thus, Thcorem 4.3 demonstrates that even after attaining 100% accuracy D Drrains the model continues to exhibit declining output diversity.
Takeaway 3: In the absence of KL divergence, output diversity continuously declines.
This diversity-com collapse phenomenon has been reported in the literature (Cui et al. 2025) and can impair model's 's ability to generalize. To address it, many techniques have been proposed, the most common being KL regularization. To better understand its role, we analyze the stable point of the model under KI. regularization, highlighting both its advantages and limitations.

---

Published as a conference paper at ICLR 2026
PED 280
(a) Test Accuracy
(b) Train Accuracy
(c) Output Diversity
(d) Influence of KL
Empirical results of PG training. Both PG and continual SFT are initialized from the same base model. Figures (a)-(c) illustrate the training dynamics of test accuracy (under greedy de- coding), training accuracy (under temperature sampling), and response diversity (under temperature sampling). Figure (d) shows how different KL regularization strengths affect the final models.
Figure 2:
Theorem 4.4 (The effect of KL regularization). When r = p= and À >> O, the stable point of the PG model satiisfess the following, under Assumption 3.7: For any fixed i. either [(kkj)(k] = or qii.j)(k) x (i, (( exp(p(i.j)[ (A)))). Here [kk] is the probability of outcome lo in softmax (oftmax(((i.j)). j)), q [i,j) is the probability of outcome k in the base model, and p(i.j)(M is the probability of tuple i, belonging valid path given output probability {q(i, ).j.k.
This result shows that KL regularization constrains the trained model to remain close to the base model, thereby preserving some of its diversity. This effect is double-edged sword. Consider valid next node k for which the base model assigns low probability, i.e., qbase (i. [)]] is small. On the one hand, KL regularization prevents qé,, [k] from becoming arbitrarily small, increasing the chance of generating valid paths involving k. On the other hand, it also prevents qi, [)] from be- coming very large, limiting potential gains when the base model's prior is suboptimal. This tradeoff explains seemingly contradictory findings in recent literature: when the base model is already capa- ble, KL regularization preserves diversity and improves generalization, but when the base model is weak, the regularization may hinder learning by overly constraining policy updates.
Takeaway KL regularization explicitly acts as a inversit--presevvii mechanism, provided that the base model is reasonably capable, but this comes at the cost of reduced train accuracy.
4.2 EMPIRICAL VALIDATIONS
The results presented in Figure where compare PG with different KL regularization factor against continual SFT. All models are initialized from the same base model after SFT training, while continual SFT means training the model for more time steps the same SFT dataset DSPT The empirical results match the takeaways we summarized from our theoretical findings, as detailed below.
Takeaway 2: In Figure 2a as the training progresses, the test accuracy of Continual SFT con- stantly decreases, while all the PG methods can achieve improvement, since they benefit from exploration-driven training data. Takeaway 3: In Figure 2b and 2 we can see that PG without KL regularization progressively achieves and maintains 100% traminng accuracy, but its output di- versity, i.e., the average number of distinct correct paths generated over 100 sampling trials for the same source-target pair, keeps decreasing during training. the end, the model eventually produces only one path per pair. Moreover, shown in Figure 2a when the diversity diminishes, continued training degrades test accuracy. Takeaway 4: As a comparison, PG with KL regularization main- tains high output diversity in the end, but their training accuracy is limited. This trade-off is further stated in Figure 2d with higher factor X, the model can have a higher output diversity and lower training accuracy. Along with Figure 2a it is shown that KL regularization prevents the model from deviating too far from the base model in terms of both diversity and training accuracy. This mitigates overfitting but also caps potential gains in test accuracy.

---

Published as a conference paper at ICLR 2026
5 ANALYSIS OF THE SLEERRINGG-BSED PLANNING MECHANISM
In this section, we analyze the Q-learning mechanism for language-model planning under two differ- ent reward designs. We show that stepwise process rewards enable convergence, preserve diversity, and remain valid under ooo-py sampling, whereas outcome rewards collapse to trivial solutions. Our analysis begins under Assumption [3] for both reward types, and we then extend the process- reward analysis to a more concrete linear Transformer model without this assumption.
5.1 THEORETICAL ANALYSIS
To analyze the structure and convergence of the Q-learning stable point, we introduce mild assump- tion, which we call the persistent exploration assumption about the RL-based learning dynamics.
Assumption 5.1 (Persistent exploration). At training step i, let i, = Ueargets Jr. = -um kit, respec- tively. denote the target, current, and next nodes. We assume for every (t,j,k) k). EN 0 such that
ij,k
T-1 Σ T-000 T
lim inf de
Prrop ijjk
Under the persistent exploration assumption, every coordinate is updated frequently enough to allow convergence analysis. In practice, this assumption is usually satisfied, for instance:
Lemma 5.1. Training with e-exploration (i.e., exploring each alternative action with probability proportional to c) satisfies the persistent exploration assumption.
With the outcome reward, the signal merely verifies whether the entire sequence constitutes valid path ending at target i. It does not differentiate between current nodes j or candidate next nodes k when i A As a result, at stable point, all logits collapse to the same constant G2 for each fixed target causing the parameters to lose structural information, as stated in the theorem below.
Theorem 5.1 (Stable points of outcome reward). Assume the RL-training uses the outcome re- ward R(u, m) = SuEP Oumsi Uuren and a stable point exists under persistent exploration (Assump- tion 5.7 and Assumption 3.1l Then, any stable point of the Q-learning model, for each fixed target and k + i, all logits [kk, take the same value depending only on i.
With the process reward, the update rule accounts for both adjacency and target conditions. The next theorem establishes that the process converges to well-defined limits that capture the underlying graph structure.
Theorem 5.2 (Stable points of process reward). Assume Assumption 3.1 holds and the process reward is used, iee. m) = 8um+1 Muran ( Mon+ 186g- Suppose the score vector f(i,j) T is initialized at zero and updated under the persistent exploration assumption with learning rate n,. Then, in the Q-learning model,
co, f(() (i) [)
A i), and for k +
1, A[ and R[i, k] 1,
0, exactly one of (A k] = 1) or (Rli, k] = A[j, and R O. Here " "->" denotes convergence or "tend to". Moreover, the convergence is linear; the effective rate depends on 7, and the update proportions Norop-
(1.5)k
-1,
Takeaway 5: Different from PG methods, in Q-learning, relying solely on the outcome reward signal can cause Q-value bias, whereas introducing process rewards mitigates this issue.
To gain further insight in a setting closer to practice we analyze a simplified but concrete one-layer, single-head linear Transformer without the abstraction of Assumption 3.1
Assumption 5.2 (Linear transformer Wang et al. (2024b)). We work under the simplified Trans- former setting in Wang et al. (2024b): (I) The token embedding matrix and the output weight matrix are both sei to the identity; (2) Attention is fixed entirely on the target node Unrres so the attention block contributes only the value lookup WV [Vrarget, :]; (3) All layer normalizations are removed, and the feedforward block is replaced by a linear map of the form FFN(X) = X.M Under these assumptions, the logit decomposes as um41 [k] Wd [um WV [Margett ki,, where Wad arises from the feed-forward weights and WV from the value matrix of the attention block.
8

---

Published as a conference paper at ICLR 2026
Despite this simplification, the analyzed results remain consistent with the experiments of real Trans- formers, as demonstrated in Wang et al. (2024b). This formulation aligns with the actual 1-layer 1-head Transformer architecture, offering greater practical utility compared to the abstract function [kk] The subsequent result characterizes the set of stable points under this decomposition and consistent with the structural limits established in Theorem (5.2
Theorem 5.3 (Stable points of process reward), Assume Assumption 5.2 holds. For linear trans- former, assume training uses the process reward, and the persistent exploration condition holds. At a stable point of the Q-learning model, for each k there exists Cj E IR. such that
WMl [j,
A[j, 1 Ck2 WV (i, k) R(t, Cg- M WV) is a stable point. Hence, the set of stable points ((WM, WV):
Conversely, any such (WM, Ch E IR & E [|]]]]
Theorem 5.T shows that only outcome reward used, the learned logits collapse to a con- stant across all states for given target. In contrast, Theorem [5.2] and Theorem [.3 show that with persistent exploration, process rewards can preserve adjacency and reachability (note that in Theorem 5.3 the constant Ck is immaterial in terms of path planning, since for any stable point, üm+11
[Mm ki] + W (urrges, k) A k] + R(i, k] - 1, which is the same as Theo- Moreover, the convergence holds even under off-pllicy sampling, and action diversity is preserved because all feasible next nodes converge to the logit value 1.
k
rem 5.2
Takeaway 6: Compared to PG methods, Q-learning can operate off-policy and better main- tains output diversity.
5.2 EMPIRICAL VALIDATIONS
We first examine the training and test accuracy results in Figure Sal where we compare Q-learning under different reward designs and sampling policies. All models are initialized from the same base model. Figure 3b states the aiccurcty-accuracy trade-off of Q-learning models, policy gradient models, and the continual SFT model (under different temperatures). Figure 3 illlsttrates the logits of an on-policy Q-learning model with process rewards and fixed attention (attention fixed on the target node utargat). The model is initialized from the same base model as in Figure 3 and is further trained with reinforcement steps on all pairs (s, t) E DSFT where t E [20]. In each row i, plot the logits for nodes 0-20 (normalized to [o, 1)) when the current node is 0 and the target node i. White indicates larger logits, black indicates smaller logits, and green frames highlight nodes that are both children of node 0 and ancestors of i, corresponding to valid outputs. The empirical results consistent with the takeaways introduced above, detailed below.
Takeaway 5: In Figure 3a Q-learning with process rewards achieves comparable training accuracy and significantly better test accuracy than the PG model, while Q-learning with outcome rewards collapses and converges to near-zero accuracy on both training and test sets, Examining each row of Figure 4 we observe that the logits of feasible nodes gradually increase and converge to the largest values within their respective rows, which aligns with Theorem and
confirms that process rewards enable the model to recover the correct graph structure. Takeaway In Figure [3a] off- policy Q-learning with process rewards attains training and test accuracy comparable to on-policy O-lcarning with process rewards, demonstrating that Q-learning can operate off off-policy. Finally, Figure 3b further highlights that the Q-learning process rewards preserve output diversity. Figure also reflects this phenomenon: within each row, the logits of feasible nodes become increasingly close one another (approaching white) over time, indicating convergence to diverse but correct transitions.
6 CONCLUSION
In this paper, we analyze the benefits and limitations of reinforcement learning in language model planning through the lens of learning dynamics. Our theoretical analysis shows that supervised fine-tuning introduces spurious co-occurrence solutions, while policy gradient and Q-learning out- perform SFT primarily through exploration. We further identify critical drawback of basic pol- icy gradient-diversinn collapse-and and show that Q-learning mitigates this issue while supporting off- policy learning. These insights clarify the mechanisms behind the recent success of RI RL-based

---

Published as a conference paper at ICLR 2026
AGDURO PM
Phii culari Me Prosee
ado
(a) Train Accuracy and Test Accuracy of Q-Learning
(b) Output Diversity vs. Accuracy
Figure 3: Empirical comparison between Q-learning and PG. Figure (a) shows the training dynamics of training and test accuracy (under greedy decoding). Figure (b) compares the Pareto frontiers of output diversity and accuracy on the training and test sets (under temperature decoding).
(a) Epoch 10000
(h) Epoch 30000
(c) Epoch 100000
(d) Epoch 300000
Figure 4: Heatmap normalized logits from the Q-learning model with process reward. For each row i, green blocks indicate valid next nodes given the current node 0 and target node i. The logits corresponding to these valid actions consistently increase during training.
approaches and highlight principled research directions, such as leveraging Q-learning for robust, scalable, and generalizable planning in language models.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We attach all the source code necessary to reproduce our experimental results in the supplemen- tary materials. As for the theoretical results, complete proofs of all the theorems are provided in Appendix C D and B for full transparency and verification.
REFERENCES
Ziwei Chai, Tianjie Zhang, Liang Wu, Kaiqiao Han, Xiaohai Hu, Xuanwen Huang, and Yang Yang. Graphllm: Boosting graph reasoning ability of large language model. arXiv preprint 10X0v:2310.05855, 2023.
Nuo Chen, Yuhan Li, Jianheng Tang, and Jia Li. Graphwiz: An instruction-folloiing language model for graph computational problems. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 353-364, 2024a.
Runjin Chen, Tong Zhao, Ajay Jaiswal, Neil Shah, and Zhangyang Wang. Llaga: Large language and graph assistant. arXiv preprint arXiv:2402.08/70 2024b.
Tianzhe Chu, Yuexiang Zhai, Jihan Yang, Shengbang Tong, Saining Xie, Dale Schuurmans, Quoc V Le, Sergey Levine, and Yi Ma. Sft memorizes, rl generalizes: A comparative study of foundation model post-training. arXiv preprint urXiv:2501.17161, 2025.
Andrew Cohen, Andrey Gromov, Kaiyu Yang, and Yuandong Tian. Spectral journey: How trans- formers predict the shortest path. arXiv preprint arXiv:2502.08790 2025.
Ganqu Cui, Yuchen Zhang, Jiacheng Chen, Lifan Yuan, Zhi Wang. Yuxin Zuo. Haozhan Li, Yuchen Fan, Huayu Chen, Weize Chen, et al. The entropy mechanism of reinforcement learning for reasoning language models. arXiv preprint arXiv:2505.22617, 2025.
Xinnan Dai, Haohao Qu, Yifen Shen, Bohang Zhang, Qihao Wen, Wenqi Fan, Dongsheng Li, Jiliang Tang, and Caihua Shan. How do large language models understand graph patterns? a benchmark for graph pattern comprehension. arXiv preprint arXiv:2410.05296 2024.
Xinnan Dai, Kai Yang, Jay Revolinsky, Kai Guo, Aoran Wang, Bohang Zhang, and Jiliang Tang. From sequence to structure: Uncovering substructure reasoning in transformers. arXiv preprint 20445 2025.
Murtaza Dalal, Tarun Chiruvolu, Devendra Chaplot, and Ruslan Salakhutdinov. Plan-seq- learn: Language model guided rl for solving long horizon robotics tasks. arXiv preprint arXiv:2405.01534, 2024.
Artur Back De Luca and Kimon Fountoulakis. Simulation of graph algorithms with looped trans- formers. arXiv preprint arXiv:2402.01107, 2024.
Jiayan Guo, Lun Du, and Hengyu Liu. Gpt4graph: Can large language models understand graph structured data? empirical evaluation and benchmarking. arXiv preprint arXiv:2305. 15066 2023.
Xiaojun Guo, Ang Li, Yifei Wang, Stefanie Jegelka, and Yisen Wang. G1: Teaching Ilms to reason graphs with reinforcement learning, arXiv preprint arXiv:2505.18499 2025.
Subbarao Kambhampati, Karthik Valmeekam, Lin Guan, Mudit Verma, Kaya Stechly. Siddhant Bhambri, Lucas Paul Saldyt, and Anil Murthy. Position: LLMs can't 't plan, but can help planning in LoM-uodulo frameworks. In Forty-first International Conference on Machine Learning, 2024.
URL https::/ wtt.://wwww..eettteteeeoeo htte//ooummmmiieee (aerrummid-hhhhBJ EmH4z
Xufang Luo, Yuge Zhang, Zhiyuan He, Zilong Wang, Siyun Zhao Dongsheng Li, Luna K Qiu, and Yuqing Yang. Agent lightning: Train any ai agents with reinforcement learning, arXiv preprint ar5iv:2508. 03680, 2025.
11

---

Published as a conference paper at ICLR 2026
Zihan Luo, Xiran Song, Hong Huang, Jianxun Lian, Chenhao Zhang, Jinqi Jiang, and Xing Xie. Graphinstruct: Empowering large language models with graph understanding and reasoning ca- pability. arXiv preprint arXiv:2403.0m 2403.04483, 2024.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5600 2013.
Ida Momennejad, Hosein Hasanbeig, Felipe Vieira Frujeri, Hiteshi Sharma, Nebojsa Jojie, Hamid Palangi, Robert Ness, and Jonathan Larson. Evaluating cognitive maps and planning in large language models with CogEval. Advances in Neural Information Processing Systems, 36, 2023.
Nanda Neel, Chan Lawrence, Lieherum Tom, Smith Jess, and Steinhardt Jacob. Progress measures for grokking via mechanistic interpretability. In International Conference on Learning Represen- tations, 2023.
Bryan Perozzi, Bahare Fatemi, Dustin Zelle, Anton Tsitsulin, Mehran Kazemi, Rami Al-Rfou, and Jonathan Halcrow. Let your graph do the talking: Encoding structured data for Ilms. arXiv preprint arXiv:2402.05862. 2024.
Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Halcrow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. Advances in Neural Information Processing Systems, 37:78320-78370, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Amrith Setlur, Nived Rajaraman, Sergey Levine, and Aviral Kumar. Scaling test-time compute without verification or rl is suboptimal. arXiv preprint arXiv:2502.12118, 2025.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathemati- cal reasoning in open language models. arXiv preprint arXiv:2402.033300 2024.
Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu, and Yueting Zhuang. Hug- gingGPT: Solving AI tasks with ChatGPT and its friends in Huggingface. Advances in Neural Information Processing Systems, 36, 2023.
Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: A flexible and efficient rlhf framework. arXiv preprint arXiv: 2409. 19256, 2024.
Trieu H Trinh, Yuhuai Wu, Quoe V Le, He He, and Thang Luong. Solving Olympiad geometry without human demonstrations. Nature, 625(7995):476 482, 2024.
Karthik Valmeekam, Mattew Marquez, Alberto Olmo, Sarath Sreedharan, and Subbarao Kambhampati. Planbench: An extensible benchmark for evaluating large language models planning and reasoning about change. In A. Oh, T. Naumann, A. Glober- son, K. Saenko. M. Hardt, and S. Levine (eds.), Advances Neural Information Pro- cessing Systems, volume 36, pp. 38975-38987. Curran Associates, Inc., 2023a. URL cctppppp.ccppppppe 7 (e922ccedeescc//raadde2222255553 Paper Datasets and Benchmarks pdf
Karthik Valmeekam, Matthew Marquez, Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models-a critical investigation. Advances in Neural Informa- lion Processing Systems, 36, 2023b.
Guanzhi Wang, Yugi Xie. Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint 23X 2023a.
12

---

Published as a conference paper at ICLR 2026
Heng Wang, Shangbin Feng, Tianxing He, Zhaoxuan Tan, Xiaochuang Han, and Yulia Tsvetkov. Can language models solve graph problems in natural language? Advances in Neural Information Processing Systems, 36, 2023b.
Lei Wang, Chen Ma, Xueyang Feng, Zeyu Zhang, Hao Yang. Jingsen Zhang, Zhiyuan Chen, Jiakai Tang, Xu Chen, Yankai Lin, et al. A survey on large language model based autonomous agents. Frontiers of Computer Science, 18(6):1-26. 2024a.
Siwei Wang, Yifei Shen, Shi Feng. Haoran Sun, Shang-Hus Teng, and Wei Chen, Alpine: Unveil- ing the planning capability of autoregressive learning in language models. Advances in neural information processing systems, 37: 19662-119888, 2024b.
Qinzhuo Wu, Wei Liu, Jian Luan, and Bin Wang. Toolplanner: A tool augmented Ilm for multi gran- ularity instructions with path planning and feedback. arXiv preprint arXhv:2409.14826 2024a.
Xixi Wu, Yifei Shen, Caihua Shan, Kaitao Song. Siwei Wang, Bohang Zhang, Jiarui Feng, Hong Cheng, Wei Chen, Yun Xiong, et al. Can graph learning improve planning in Ilm-based agents? Advances in Neural Information Processing Systems, 37:5338 5383, 2024b.
Jingkang Yang, Yuhao Dong, Shuai Liu, Li, Ziyue Wang, Haoran Tan, Chencheng Jiang, Jiamu Kang, Yuanhan Zhang, Kaiyang Zhou, al. Octopus: Embodied vision-language programmer from environmental feedback. In European conference on computer vision, pp. 20-38. Springer, 2024.
Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Shiji Song, and Gao Huang. Does inforcement learning really incentivize reasoning capacity in Ilms beyond the base model? arXiv preprint (tt.:2504.13337 2025.
Guibin Zhang, Hejia Geng, Xiaohang Yu, Zhenfei Yin, Zaibin Zhang, Zelin Tan, Heng Zhou, Zhongzhi Li, Xiangyuan Xue, Yijiang Li, et al. The landscape of agentic reinforcement learning for lms: A survey. arXiv preprint 00Xiv:2509.02547, 2025a.
Kaiyan Zhang. Yuxin Zuo, Bingxiang He, Youbang Sun, Runze Liu, Che Jiang, Yuchen Fan, Kai Tian, Guoli Jia, Pengfei Li, et al. A survey of reinforcement learning for large reasoning models. arXiv preprint wxxiv:509900822 2025b.
Hanlin Zhu, Baihe Huang, Shaolun Zhang. Michael Jordan, Jiantao Jiao, Yuandong Tian, and Stu- J Russell. Towards theoretical understanding of the rrversal curse' via training dynamics. Advances in Neural Information Processing Systems, 37:90473-90513, 2024.
13

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS
In this paper, the core conceptual framework and its iterative development were driven by human researchers. LLMs served strictly a supporting capacity, primarily employed for linguistic re- finement of the manuscript enhance readability while preserving original technical content. The whole paper is carefully supervised, reviewed, and modified by the authors who maintain complete responsibility for the scientific validity, technical accuracy, and ethical integrity of this work.
B MORE RELATED WORKS
B. PLANNING OF LLMS
Planning is fundamental component of human intelligence and autonomous agents. Several stud- ies have evaluated the planning capabilities of LLMs trained without reinforcement learning, such as CogEval (Momennejad al. 2023) and Blockworlds (Valmeekam al. 2023b). These works consistently report negative results, suggesting that LLMs lack inherent planning abilities. In con- trast, models such as o/ show the ability to solve such problems, though the underlying mechanisms remain unclear.
On the other hand, LLM-based agents have demonstrated remarkable competence performing real-world planning tasks, even without RL training (Wang et al.) 2024a). Many of these plan- ning tasks can naturally abstracted as path planning problems on graph. For example, in tool-augmented agents (Shen et al. 2023), tool dependencies can be modeled as a graph where nodes represent tools and edges represent dependency relations (Wu et al., (024b). Planning, in this context, involves finding path of tools to fulfill the user's request. Similarly, in mathemati- cal reasoning agents (Trinh et al. 2024). theorem dependencies form graph where constructing a proof is equivalent finding path. In game-playing agents such Voyager (Wang et al. (2023a), skill dependencies create a graph structure where planning determines the sequence of skills needed accomplish tasks. These observations motivate abstraction of planning path planning problem in this work.
Agents trained without RL face key challenges: (1) supervised fine-tuning loss is misaligned with the agent's ultimate objectives, and (2) real-world data is scarce. RL addresses the first issue by explicitly optimizing for the end goal through reward signal, and the second by generating exploratory data. Consequently, RL significantly mitigates these limitations and improves perfor- mance (Zhang et al. (2025a). Our paper further examines the benefits of RL over SFT, well the limitations of RL, providing insights for future research directions.
B.2 RL FOR LLMS
Recently, RL. has been widely adopted to enhance reasoning capabilities in language models, exem- plified by milestone systems such as OpenAl's ol and DeepSeek-RI. This paradigm has inspired new wave of reasoning-focut 2025b). State-of-the-art LLM-based agents also commonly employ RL (Zhang et al. (025a).
models, including Qwen-3 and Phi-4 Reasoning (Zhang et al.,
Despite its empirical success, the mechanisms by which RL improves LLM performance remain an active area of research, with current understanding scattered across multiple works. For instance, Chu et al. (2025) empirically compares SFT and RL on reasoning benchmarks, concluding that RL provides better generalization. Theoretical analysis in (Setlur et al.] (025) further shows that any verification-fies approach, such as SFT, is suboptimal. Additionally, Yue et al. (2025) identifies an entropy mechanism, establishing and empirically validating trade-off between entropy and accu- racy during RL training.
In this paper, we focus path planning as case study and derive results consistent with prior work: (1) SFT tends to memorize training data and produce co-occurrence-dee outputs; (2) RL surpasses SFT primarily through exploration; and (3) diversity collapse occurs during PG training. Beyond these findings, we uncover evidence suggesting that Q-leanning may offer advantages over policy gradient methods, introducing new perspective on RI. for LLMSs.
14

---

Published as a conference paper at ICLR 2026
B.3 GRAPH PROBLEMS WITH LANGUAGE MODELS
Graph problems serve as valuable testbed for analyzing the reasoning capabilities of language models. From an empirical standpoint, several benchmarks have been proposed (Guo et al. 2023
Wang et al, 2023b Dai et al., 2024), spanning a spectrum of tasks: classic graph problems (e.g., connectivity, math-finding. and pattern detection), graph neural network (GNN) benchmarks (e.g., node and graph classification), and semantic graph-based question answering (e.g., on knowledge graphs). Without additional training, LLMs generally underperform
these tasks. To improve
performance, early approaches leverage instruction tuning and DPO (Luo et al. 2024 Chen et al.
2024a; Perozzi al. 2024 Chai et al. 2023 Chen al. 2024b), while later methods employ
RL (Guo et al. 2025). which consistently achieves superior results.
There three major paradigms for analyzing how transformers solve graph-related reasoning tasks. The first is mechanistic interpretability. which reverse-engineers trained model weights (Neel et al,, 2023 For example, CChhnn et al. (2025) observed that transformers implement spectral algorithm to compute shortest paths. However, this paradigm largely relies on empirical observation without theoretical grounding. The second paradigm based expressiveness analysis (Dai et al. 2024
Sanford et al. 2024, Dai et al. 2025 De Luca & Fountoulakis 2024), constructing weight configu- rations that enable transformers simulate algorithms. Yet. such configurations are often unrealistic for transformers trained via SGD (c.g., embedding vectors explicitly set to 1.2..... L (Dai et al.) 2.24)). The third paradigm investigates gradient dynamics, which both practical and challenging due to the non-convexity of the optimization landscape. Prior work has analyzed aath-finding in directed graphs (Wang et al. 2024b) and compositionality paths (Zhu et al. 2024).
To the best of our knowledge, this work presents the first analysis of RL gradient dynamics in LLMs. Our results explain why RL-based methods outperform SFT approaches and highlight the potential advantages of Q-learning-deee -driinn methods, opening promising directions for future research.
C APPENDIX FOR SFT
C.1 PATH PLANNING ALGORITHM IN TRANSFORMER
Algorithm A handcrafted path planning algorithm
1: Input: Adjacency matrix A. reachability matrix R, source node s, target node and set current node
2: Set path p= [s
3: while * /. do
4: Obtain
{k: A(,
and R((,)) 1}
5: Randomly sample next node & from S
6: Append to path P, and set =
7:
end while
8: output path
C.2 PROOF THEOREM 3.1
Proof. The next-token prediction cross-entropy loss can be written as
i t=
Σ
Σ 6k 6k=um++ log um [k)]
WEET m21
Under the assumption that the output distribution ún depends only on the target node Utarget and the current node Um. we can aggregate identical terms, and express the loss as
N Nugger Wuggees Moye
exp(f (Maageee Uh [k]] Σv exp(f (Murgett um [k-))
Σ
Σ A
N Murrett chem Je'
Σ
log
Σx N Musses Jkm
15

---

Published as a conference paper at ICLR 2026
If Ex um F O, the expression in brackets is the cross-entropy between the empirical distribution and the distribution of doing softmax on vector f [Muarget Mmm) It is minimized when
Uuprges Unarget Um [k])
N Naunga Men.k
Σ exp(f( (uurget Um [k')
ke Mhorge Mem di
Ex N Mange Man,
0, the loss does not depend on Uuureet Um) so it can be any valid probability
distribution.
D APPENDIX FOR POLICY GRADIENT
D.I PROOF OF THEBRREM
Proof. In this case, the loss function of policy gradient is
t= Σ dUEP
log üm [mm+a]
Σ
log úm [ummtal)
DALA
m21
DEDRLINP
m21
D.2 PROOF OF THEOREM 4.2
Proof. We first rewrite the loss function as
8
Σ
Σ m21
log úm [um+1]
R.
R.P.t bii
(-log Ex exp ((((())
exp(f(i.))]
DEDRIiNP
where N.P.P.t denote the number of times that Mnaget Um j and um+1 k for m 1 in DRLt n.P Then we can take the gradient and get
Oe
R.P.P. ijk +
exp(f(i,
Σ
N.P.P. Lik
(f
Σn exp f((T
For a wrong tuple j, (where is not adjacent with j or cannot reach i), N R.P.t. kJ.k
always zero
Thus, the gradient always positive. On the other hand, we will also have
Oe
E. exp exp(((i,,)) Σn exp(f(1.j)
Σ
Σ
N R.P. ijik
+
Σ
R.P.P. GJ
0.
D.3
PROOF OF THEOREM 4.3
Proof. Note that
[((i (llg [C(i.)) log exp(f(i.jj[k))
KL(UCO.)) softmax(((i..))
Σ
ENeO(H exp(f(i,,))))))
kEC
exp(((i,,))) k))
log ((i,)))
log
[((, kEC()))
Σ ExeCOs exp (xp(f(i,j)kk]]
Thus, it is sufficient to prove
log exp(f'(1.
E log exp(f++1 (b))
Eveo(n exp(( (k)j)
Exeonn exp [(+t((
kEC(
AECH)))
16

---

Published as a conference paper at ICLR 2026
According to the gradient, we have that
px (1.5))
(1,5))
N.R.P. -Nij.
exp(f(i,))))) + ExECOD exp(((i.)) MEC(L))
E
NRP.P. Nish
Here 7i is the step size.
Let N R.P.t
Exect,)) N.P.P. R.P.P. 'j.k'' then due the on-policy updating in policy gradient, know is the counter of outcome k for Ni. independent multi-nomial random variables with This means that
Nij
that N R.Pt bigh
(x.)
parameters
ENecu, exp(P
keC()))
B[ (i, (A.))
Moreover, since log (ExECtU.)) (xp is a convex function, we also have
E log
Σ
Exp(P*+((1,
log
Σ
(EC(()))
WEO(K))
log
Σ exp(f (2.1)
WEC((1))
Because of this, we have
exp(F'(1,))
exp(F'+1 (k,))
Σ log E EMEc(s) exp(f' (xp((((i,
E
log
E Ei'EO(i)) exp (ft+1 (1, (AT)
kEC(L,)
kEO()
Σ
(1.D) - E
Σ (EC())))
p (i,5)))
- -C((,)))
Σ exp((((i.))) (i.
(EC(L.)
log
Σ
exp(rt+) ((.0)K1) (x.()++((.)))11)))
(EC()
IV
0.
D.4 PROOF OF THEOREM 4.4
Proof. When O, the loss function is
{ log q(i.))(A
(i. (111))111111 }).
8= Σ R.P.P.t (- ijik
log (b)) + A
Σ
R. biie
log (ogg
Lo.k
where N.R. R,t denote the number of times that Ugarggt i, Um DRL.
and um+1 k for m in set
We can take the gradient and get:
0e
R.P. -N..PP. [k. Σ
Nijk A.j (1- log q buse (i,
[kk
[k}
Rtt Nijkk
(Kx))) log
q buse (k.
17

---

Published as a conference paper at ICLR 2026
Taking expectation, get:
al
-E bick
(k) Σ
B[N.PPP AE
[().))( log g (z. [k)
[R]
R. E[N ( log (kk
Σ ki
q base (i, (i,j)[k)
Letting N
Ex
Rut Like
then due on-policy training, we have that E[N [A).
R.t ((.))) and E[N R.
Hence
de (f
-N + ((.))())) Σ [d(i:)lWpp(iii)
[kk]
qii,j)
+AN q(i.))( {(i.)((())) -q(1.)) log quuss [1,
AN 1.
q(..)) log h (2,
(i, [A]]
N ((. Σ ((.5)) (D(,)))M) p(i,5)(k] À log
À log base [k]]
base
[k)
The stable point must satisfy that, for any tuple j, k E
[[()) = 0. And we claim that in this
at
case, for fixed j and any k such that qi., [k]] > 0, their [k] + A log q(i,)( should equal. Otherwise
q(.)) always look for k* arg ((.))) ((.,))[A] Alog q (L
and its expected gradient E
at
is strict negative.
E
APPENDIX FOR Q-LEARNING
E. PROOF LEMMA 5.1
Proof. Fix any triple (i. j. k). Consider training sequences whose first two nodes satisfy Usouree E V and Uaurge By the definition of the training process, P Piuyource E V, Munget = i) 0. Under c-exploration uniform over V,
VEEV
P(next node v) 2 c/1//
Condition on the event (Msourcee E V, Utarget
= i}. Then in the next two decisions,
P(u2 = j Msource E V, Utarget
i) 2 e/ IP( = k Usnurce E V, Unarget i, 12 = j) 2 e/
Hence
P(uuarge = i, U2 = j. U3 k) 2 Po(e/|V1)2 0.
Each occurrence of (Mangett, U2, Wz) (i, j. k)) triggers one update of [k)) Since each occur-
rence yields an update of [})), we obtain
lim inf T00 T
Σ 8(1..)
Nprop- >
which is the persistent exploration condition.
E.2 PROOF OF THEOREM 5.1
Proof. At stable point, the expected update of each coordinate vanishes. The per-step loss is
f (f(Muuree Un [um+1] Sugp 6 60m4 Euwgn
{max (Munger Um+1 [kk?))
18

---

Published as a conference paper at ICLR 2026
Taking the gradient with respect to f( (Murget Um) [um++1] and setting the expectation to zero yields, for every triple (i,j,k) k),
f(i, Blo d + max [K]]]
of i, the expectation term equation vanishes, so
[( max f(i,k)
which does not depend on j. Thus, for each o we have
( [max f(i, (KK"), R(i, f(i. [b. max f(i, [)()) max
max max f(i,, k'([""- #s
ke أن kn
Here the second equality separates the case K' = i from k' 1 i: if i, the chain terminates immediately, corresponding to the - - path, so we take the value f(i. ki directly. If o we must expand one step further, leading max4" f(i,, [(),") The last equality then observes that including the = term does not change the maximization, since is already dominated by the expansion over k' * i. Thus the final expression simply enumerates all possible terms with the restriction k' * i. This expression no longer depends on or j.
Therefore, for each fixed and + i, all [ik] take common value c;, independent of j and k.
E.3 PROOF OF THEOREM 5.2
Proof. For clarity, we introduce two notations that will be used repeatedly. First, for i, k E [nd] and iteration t, define
max p( (2,
Second, we write E Anc(i) if there exists m 2 such that (A (Am)k, i] = 1, i.e. & an ancestor of i.
The per-step loss under the process reward is
t= (f( Murget, Um [1 (5 (6um-1
δ (M +1)))
{max fMuurgat um++ (KT]) ke'
Taking the gradient with respect to the active coordinate ['uarrers Um [1m=- 1] gives
0e
( (Uuaget, Umm [Umn++] 1] -: -uupe + (00m Mmm
max f (Muagett 1) by
Of (uureet mm
[1
Applying gradient descent with learning rate 7) yields
f( (uugges Um [um++] +- (1-21) ('uagea Un 11 (oum 1mm-1 (mm Um 1 + max f' fuuuret Un+ 1)
Renaming (i,j,k) (Murget, Uyn Um and writing s(1) max (i, the recursion is
p(++")( = (1 -1-200 (1 (1).))()) + 27) + + - 11 $!2)
When l: = i, by convention s(4 = 0. so the update is
1(++1) (i,j)[ = (1 - 2n) (1, (1) + 2m A[j.j].
This linear recursion has fixed point A [j. i] and solution
(t.)) = (1 - (1-2)))) A[j d].
which converges to if A[j, = and to 0 otherwise. The contraction factor is | 2ny]- For l o the limit of i must be analyzed. If ls e Ane(i) \ {i), then either & is parent of i, in which case p(t) k) - and hence S(O) - k has child with E Ane(i).
19

---

Published as a conference paper at ICLR 2026
(tt) Inductively 507 - and then (+) implies (O (i,)))) Sik 1. If k. & Anc(i), then all children of l also satisfy r & Anc(t), and inductively S + 0, giving fit (i,b))r) =
(1-2n) (i. (i. (r) + Thus sit O. Therefore the limit / 1ifk (nc and $ + 0 otherwise.
For + i, substituting the limiting $12 into (+) gives
2n) (1. = (1-2n) (i.J)[a] + (n)( - 1) + s(
If A[j k| 1 and & E Anc(i), then s/1 L, so (i. + 1. If A k] and k & Anc(i), then 1 + 0, so (t..) + O. [j.
and Anc( then / - 1, so the recursion (1 - 20) (i. [A]] implying (1. + 0. If A [j. k] = 0 and k & Anc(i), then 5/0 + 0, so the recursion 1(++1)(i, (i. [A] = (1 - 2n) [i. 27y, which converges to -1. These limits match the cases in the theorem.
is (i.))
To establish rates, define the weight error E\ (i,j,k) = (i. - f [K]] and the max error (i, 5 ST When i, the recursion is
W e++1 (i, j. i) = (1- 2n) et (i,j,))
so each update contracts the error by [1 2n). Under persistent exploration, the coordinate (f. d)
is updated with positive frequency Noppo so in global time
|e," (i.5. S C (I1 2nl
NP CC
for any e > O and large enough
When k + i, the recursion is
W W. (i.j), = (1 - 2nj e (1.j,k) + 23 2 The error ef k) depends only on {er" (i, k r) : r child of k:). Along any directed path == "D
+ Um = i, the enror at h can decay only after the error at v has already decayed, and so on. Thus, the effective contraction factor for et (i. k) is the product of the per- edge contraction rates
II ==
(11 2ny
prop vni ).
Formally, by induction, for any & > 0 and sufficiently large t we have
1e," (i,j,A))
max palhs pok-ti
II
1 21y
Narop #n41
(Mmsn++1)Ep
Therefore, all iterates converge Minnerrgeinnearly in global time, with effective rates determined jointly by n and the update proportions Nireer This completes the proof.
E.4 PROOF OF THEOREM 5.3
Proof. Let i.j.k denote the asymptotic proportion of triples (Margets um: Mmx+1) = (i. k) oc-
curring in the generated sequences at the stable point, under the given sampling method and the persistent exploration condition. Equivalently, N... is the limiting frequency with which state transitions to k with target in the trajectories sampled by the model. By definition Na > 0 for all k.
At a stable point of the updates, the expected gradient with respect to each parameter must vanish. Averaging the stationarity conditions with weights N yields, for all i.j.
Σ
Npropp ij..k
W
M
[i
k]
+
W
[i, k]
Alj. k] + 1 8azk max W
[ks AT] + W (i,AA) O,
Σ
W
Li kj
W [i - A[A + 1 8a= max W M [k. A3] W/ (i,,AT)) O.
20

---

Published as a conference paper at ICLR 2026
Introduce centered variables J.A= w j. k]
1,
TAA w [i, ò,= max W (ax(W (i,
so that normalizing each sum by its positive denominator gives the block system
S + P4 Tx =
Tx + Q+ 00
(5)
where S = (S)))))E[m) Tx = (T jmj) and
Nprop
Norop
(P Ex p
(Qn)
Eg
N
Since Nprop O, every entry of Pa, Qu is strictly positive, and both are row-stochastic. Hence Qa and QuP strictly positive stochastic matrices. By the Perron-Frobenne theorem, both have a simple eigenvalue with eigenvector 1, and all other eigenvalues satisfy |^] < Thus
ker(I - PR Qa) span(1},
ker(I Q2 Pk) span
From equation 5 eliminating Tş yields S (P&Q+)S, so SR c.1 for some ck R. and then Ta = - Q 1. Returning to the definitions,
W [, A A k]
+ Ck:
W [i, k) = di= + max
(W [k. k') W W [i, Ck-
&
Substituting W M [k. KJ A(k W - + O and writing Vi == w [t, k1 a gives
Vi = 6;=b +
max A[k,(k]
V.
On a DAG, the unique {0, 1} solution of this recursion is the reachability indicator R. An induc- tion over topological order shows Vik R4, for all i, k. Therefore
w W k]
R [i, k) Ck.
Finally, note that if (Sx, T&) solves equation then does (S& + cl, Tx cl)
cl) for any c E R, since P = Q = Hence, the solution set for each & is exactly one-dimensional affine line parametrized by Ck-
Conversely, if (W-M, W is of the above form, then plugging it into the update equations shows that the expected increment is identically zero: both sides of the gradient equations cancel by con- struction, so the point is stationary. Therefore, these conditions are not only necessary but also sufficient for stability.
F EQUIVALENCE OF UNCLIPPED PPO AND POLICY GRADIENT
For a sequence u, the policy gradient objective is
(u) Σ
I((u) log úm [um+1])
m21
where R(u) r duep + Taking the gradient gives
Volpa(u) = Σ
R(u) log úm [um++
m21
Voun [um+1]]
Σ
R(u)
úm [Mm+1]
m21
For unclipped PPO, ratio between new and probabilities is formed, with the denominator detached. The loss is
EPPO-nc(u) = Σ R(u) üm [um+1]
{ūm [um+1]]
m21
Since the denominator { (úm [um+i]} is treated as constant, its gradient vanishes. Thus
Σ R(u) [um+]
v Volppo- (u)
[ūm[um+t] [um+1]]
Comparing with the policy gradient expression, the two gradients coincide. Therefore, for any fixed sequence u. unclipped PPO with a stop-gradient denominator exactly equivalent to vanilla policy gradient.
21

---

Published as a conference paper at ICLR 2026
(a) Attention weights in SFT after different training iterations
(b) Attention weights in PG (A 0) after different training iterations
(c) Attention weights in Q-leaming (Process Reward) after different training iterations
Figure 5: Empirical validation that the trained one-layer one-head transformer acts as a function of the target node and the current node. The visualization of attention maps across SFT, PG, and Q-learning training shows consistent, strong focus on the target node (token position 1).
G ADDITIONAL EXPERIMENTAL RESULTS
G.1 VALIDATION OF LEARNED ATTENTION
This subsection visualizes the evolution of attention maps during training for SFT, PG, and Q- learning. Our analysis first focuses on a one-layer, one-head transformer at various training steps (Figure 5). For each model, we compute the average attention weight over the SFT training dataset DSPT
Specifically, the weight position i in row i (where i k) represents the average attention the model assigns to the i-th token when predicting the (k: + 1)-th token, averaged over all paths in DSFT longer than k + 1. Since the underlying graph has 100 sparse nodes, path lengths in DSFT generally short; consequently, we display only the first 8 rows.
Our visualizations reveal that during SFT, the transformer allocates most attention to the target node. Combined with the residual connections, which allow access to the current node's 's information, this suggests that the model learns to predict the next node based primarily on the target and current nodes. This empirical finding aligns with the results of Wang et al. (2024b). Another interesting phenomenon in SFT that the attention weight on the target node quickly peaks and then gradually decreases, while remaining dominant. We hypothesize that this may be due to overfitting on DSPT leading the model to develop auxiliary prediction strategies. This overfitting could also explain the decreasing generalization performance of SFT observed in our experiments.
In contrast, for both PG and Q-learning, the attention on the target node increases throughout train- ing. In Q-learning, the final attention weight the target node exceeds 95%, making it the closest to the conditions outlined in Assumption 3.1
To further validate our findings, we extend the analysis to two-layer, one-head transformer. As shown in 6 which displays the attention map averaged over DSFT both layers predominantly attend to the target and current nodes. This pattern strongly supports our assumption that the trans- former operates as a function of the target and current nodes, confirming the consistency of this behavior across model architectures.
22

---

Published as a conference paper at ICLR 2026
Attention weights in SFT after different training iterations
(b) Attention weights in PG (A = 0) after different training iterations
(c) Attention weights in Q-learning (Process Reward) after different training iterations
Figure 6: Empirical validation that the trained two-layer one-head transformer acts function of the target and current nodes.
G.2 EXPERIMENTS ON ERDÓS-RÉNYE GRAPHS
Beyond the setup in Section 2 we conduct additional experiments to compare RL methods with SFT. The graph construction and initial SFT stage remain unchanged. After SFT, we split all reachable pairs into an RL. training set DRI.Train and an RL test set DOL. This yields four intersections:
DrrainTesin Drrain O DRI-Trains Dresin2Test ;= Dreain O DRL-TeSt: Dreat2Train := Dreat 0
DRI-Train- and DesttTess := Drest n DRL-Test
During the RL process, the model generates paths for pairs in DRI-Train and receives reward signals. The main difference between this setup and that in Section 2 is that the RL training set now contains new pairs that were unseen during SFT (Drest2team) Therefore, the initial model is not perfect on these new training pairs. Additionally, some pairs from the SFT training set are not used for RL training (Dra)n2Test,, which allows us measure the extent of forgetting. We consider the same
23

---

Published as a conference paper at ICLR 2026
1409
6575
Accurary Druacttve
Accuriiy on Drvestiee
8:33
Accifary on
0.03
0353
8223
3300
4008
5006
0090
33006
2000
6083
0008
6800
1000
2006
4693
0089
1900
10089
3000
4899
0093
1000
1009
Tainiin
Tririing Stop
Traning
Training Stop
00
0.033
00001
Figure 7: The test accuracy PG with different KL coefficients on four data splits after fine-tuning the SFT model on Trai-r All accuracies are evaluated with greedy decoding.
Arcurey on Drovcrras
Drructer Tea
Accuracy on Drowi
Accumity on
1900 12900e TRA0C 30008S 3900e Training Stop
00900 120000 I86000 2e3ice 1996ce
6500 120093 18K89 240000 30000
00830 1230ee Ledcee 200093 30003
Trarinng Stop
Traning Step
Training Stop
Leaminn Prrocess Rewa
Figure 8: The test accuracy of Q-learning on four data splits after fine-tuning the SFT model on DRI-Train- All accuracies evaluated with greedy decoding.
RL algorithms introduced in Section 2 PG and Q-learning, whose training curves are presented in Figures and respectively. All accuracies are evaluated using greedy decoding.
From Figure we observe the opposing effects of KI. regularization Dreain2Test and Drest2Train- PG without regularization (A 0) and less regularized PG (A 0.0001) achieve significantly higher accuracy on Drest2Train Stronger KL regularization hinders the model's ability to learn new pairs, which aligns with Takeaway 4: KL regularization reduces training accuracy. Conversely, PG without KL regularization (A 0) tends to overfit the training data and exhibits continual forgetting of previous knowledge learned during SFT. Results on Diest2Test further demonstrate that overly strong KL regularization can hinder PG's improvement. Among all settings, 10-4 achieves the best balance, indicating that a well-chosen KI. weight can improve generalization with minimal sacrifice in training accuracy.
Compared to PG and the performance observed in Section S Q-learning exhibits slower convergence in this setting. One possible explanation that the initial model performs poorly on the new training pairs, generating more failure cases and causing stronger "re-instantiation"
G.3 EXPERIMENTS ON GRAPH REPRESENTED FOR BLOCKSWORLD
We also run experiments on Blocksworld (Valmeekam et al. 2023a), a benchmark for evaluating LLM planning ability (Kabhaapptii et (024). The environment consists of blocks stacked on a table, and the goal is to rearrange the blocks from an initial configuration to target configuration using sequence of actions. We model this into a aath-findig task, in which each configuration is node in graph, and an edge connects two nodes if one configuration can be transformed into the other by a single valid action, such as moving a block from one stack to another. We consider Blocksworld with four blocks and construct undirected graph with 73 nodes: 24 configurations of a single stack of four blocks, 24 configurations with three blocks in one stack and one block on the table, 12 configurations with two stacks of two blocks, 12 configurations with stack of two blocks and two blocks on the table, and one configuration with blocks on the table.
Since accuracy comparison is not focus, all node pairs used for SFT training. The SFT dataset contains 50,000 paths sampled from the graph, with source and target nodes drawn uni- formly from the 73 nodes. During RL. training, the model generates paths for, and updated on, all
24

---

Published as a conference paper at ICLR 2026
node pairs. We use policy gradient and Q-learning introduced in Section After training, evaluate the learned weights using the metric Wang et al. (2024b), which measures the model's understanding of graph adjacency. As shown in Figure with fixed training data, SFT may not learn the complete adjacency very well. In contrast, both PG and Q-learning improve the learned adja- cency. In particular, Q-learning nearly recovers the complete adjacency, consistent with the results in Section 5
25


Reviews:
[{"review_id": "CXUJcq45Xq", "reviewer": "Reviewer_q3k8", "summary": "This paper theoretically analyzes why LLMs work well for planning problems. They formulate the problem as  a graph where the prompt defines a start state and a target, and then the agent gets rewarded for generating a valid trajectory through the graph from start state to target.\n\nThey provide theoretical explanations for...\n- why Policy gradient (PG) > SFT (exploration)\n- why PG results in entropy collapse\n- why KL regularization mitigates this\n- why Q-learning on outcome rewards experiences \"reward hacking\" (their word for inaccurately optimistic Q-value estimates)\n- How Q-learning with process rewards mitigates this", "questions": "Disclaimer: it is quite possible that I have misunderstood some of the key arguments in this paper. If these could be clarified (both for myself and for future readers) my score could increase significantly.\n\nQuestion 1 - could you clarify exactly what you mean by a \"graph\" in this setup? I think it was often ambiguous whether the paper is referring to a semantic graph (e.g. for a math proof, nodes are proof prefixes and an edge from A to B exists is B contains A, but with one more logically correct step added (or is it one valid next token added?)) or to a literal graph (vocabulary are nodes, graph is fully connected b/c any token can come after any other one). \n\nI'm also very confused about whether we are assuming the graph is markovian or not. In section 2.1, it says that each vocabulary item is a word. However, whether a transition from one word to the next is valid depends not just on the current word, but on all prior words in the sequence. Furthermore autoregressive transformers look at sequence prefixes not just words in isolation. So is a node a single word, or a sentence prefix? If it's a single word, I don't think this matches how LLM RL is done in practice.\n\nThis confusion further makes me question the Theorem 5.1 results. It makes sense that you would converge to equal Q-values for all next states if it is possible to select k' across all possible next nodes. However if nodes are prefixes it's not clear to me if this still holds. My mental model is a problem where the start token is \"Hello\" and the model gets reward 1 if its next token is \"World\", but reward 0 for anything else ever. If a transformer is trained on a batch of trajectories with \"Hello World\"-> 1 and [anything else]-> 0 it is not clear to me the theorem still holds. If a sequence is [random, sequence, of, words....] then there is no possible transition from that to a sequence [Hello, World], so I don't think Q-learning would ever give this sequence > 0 reward.\n\nQuestion 2 - I think some of the discussion about SFT vs PG was a bit unclear:\n- The paper states that \"SFT memorizes\" whereas \"PG generalizes\". The claim that SFT memorizes makes sense and was well supported. However, the claim that \"PG generalizes\" seems at odds with the analysis showing a loss of diversity. It seems like PG is not truly generalizing, merely memorizing a small set of paths from start state to the solution. The Fig 2a results do show stronger test performance of RL methods, but the \"Continual SFT\" setup was not described, so I'm not sure what it means.\n- There are 2 differences between SFT and PG: (a) PG has a notion of reward which SFT lacks (b) SFT is static whereas PG collects new trajectories online. The theoretical analysis focuses on the benefit of (b) but I think does not address (a).", "limitations": "", "rating": 2, "confidence": 3, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "- Paper carefully cites where it extends prior work\n- Claims are supported by experiments in Blocksworld, where they observe that KL regularization maintains output diversity during PG training and poor accuracy during Q-learning training with outcome rewards but good accuracy with process rewards.\n- Graphs are clear and have error bars.", "weaknesses": "See the \"Questions\" section for more details, but at a high level, I think there are 3 areas for improvement:\n\n- Make it clearer exactly how your theoretical analysis ties into real-world LLM setups. I think this could be made very explicit with a running example that you reference back to in each section (e.g. you could say something like \"suppose you have a dataset consisting of math proofs. For each proof in the dataset, the starting state is [....]. The target is [.....]. Each state is [.....]. There is a valid transition between states if [....]...\" Then when you introduce Q-learning you could say something like \"The Q-function takes inputs in this form [...question, math proof up until step X...]. But when we optimize it, we see reward hacking [...predicting that invalid proofs are correct?...] Also discuss whether the assumptions you make when constructing graphs for your experiments are realistic assumptions that match data distributions in real natural language. (you have a line \"a tool-use scenario can be modeled as identifying a valid call sequence within an API call graph\" -- this is a good candidate to be fleshed out into a full example).\n\n- The trajectory notation with u1 = start state, u2=target u3+ = actual trajectory was somewhat confusing to me. I'd suggest giving the start state and target different notation (e.g. s0, s_target) and then indexing the actual trajectory from 1.\n\n- See my questions in the \"Questions\" section about the mismatch between the non-markovian setup in real LLMs and the markovian Q-learning setup.\n\n- I think the phrase \"reward hacking\" is a bit misleading in this case. Typically this is used to refer to exploiting a misspecified external reward rather than having biased upward estimates of Q-value estimates within the policy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W8Rewwfblp", "reviewer": "Reviewer_hAiC", "summary": "This paper provides a theoretical comparison of different training paradigms for enabling language models to learn to plan—namely Supervised Fine-Tuning (SFT), Policy Gradient (PG), and Q-learning—within a controlled graph-based path planning setting.\n\nThe authors first show that SFT tends to memorize observed co-occurrences limiting generalization. They then demonstrate that PG improves over SFT despite having a similar objective, because its on-policy data generation introduces exploration, effectively augmenting the training set. However, they identify a diversity collapse in PG, where the model’s output variety diminishes even after achieving perfect training accuracy—a problem that KL regularization can mitigate at the cost of lower accuracy.\n\nFinally, the authors analyze Q-learning, showing that it faces convergence issues and detail a form of “reward hacking” when relying solely on final outcome rewards, but that these problems can be addressed by introducing a process rewards to supervise intermediate steps. Under this setup, Q-learning naturally preserves output diversity and supports off-policy learning—unlike PG.\n\nAll theoretical predictions are validated empirically on a small Blocksworld planning environment using a one-layer Transformer.", "questions": "- Your theoretical results rely on **Assumption 3.1**, which posits that the model's logits for the next token depend only on the current node $u_m$ and the target node $u_2$, effectively neglecting the influence of the preceding path and initial state.  \nCould you discuss how your analysis would change if this assumption were relaxed—for instance, if the model's predictions also depended on earlier nodes in the trajectory (the full prefix $u_1, \\ldots, u_m$)?\nIn particular, would the main results (e.g., Theorem 3.1’s co-occurrence characterization or the PG/Q-learning convergence properties) still hold qualitatively, or would new interaction terms or dependencies appear in the gradient dynamics?\n\n- In your formulation, RL exploration appears to operate only over paths between fixed source–target pairs sampled from $D_{\\text{Train}}$. Have you considered or analyzed the possibility of *exploring over source–target pairs themselves* (e.g., dynamically selecting or expanding the set of training pairs)? How might such goal-level exploration affect the generalization and theoretical results you present?\n\n- In Figure 2, PG ($\\lambda = 0$) achieves its highest test accuracy early in training before diversity collapse occurs, while KL regularization later limits accuracy gains. Wouldn't *early stopping* at this peak be a simpler and potentially more effective strategy than adding KL regularization, since it avoids both collapse and the regularization-induced accuracy loss? Did the authors consider this comparison?\n\n- **How specific are your findings to transformer architectures?**\nYour theoretical analysis and experiments use a one-layer Transformer as the backbone. How would the conclusions change for other model classes such as RNNs (e.g., LSTMs), state-space models (SSMs), or diffusion-based architectures? Do your results rely on properties unique to transformers (e.g., attention structure), or do they extend to other sequential models?\n\n-**In real-world applications, process-level reward information is often unavailable.**\nGiven that your analysis shows Q-learning may fail to converge under outcome-only rewards, how could this limitation be addressed when such intermediate reward signals are not accessible?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well written and clearly presents the theoretical analyses, making it easy to follow the derivations and main arguments.\n\n- The controlled graph-based path planning framework provides a clean setting to extensively study the differences between SFT, PG, and Q-learning for planning with language models.\n\n- The theoretical claims are relevant to the real world challenges of training LLM to plan.", "weaknesses": "**Lack of empirical validation at scale**  \nWhile the theoretical analyses are thorough and insightful, the empirical validation is limited to a small Blocksworld environment with a one-layer Transformer. It remains unclear how well the theoretical findings translate to larger, more complex models and real-world planning tasks. Additional experiments on larger datasets and models would strengthen the practical relevance of the results. Given the clearness of the theoretical contributions, it seems that it would be straightforward to validate them on larger models/environments.\n\n**Figure 1 description**\nCould you add more explanation to Figure 1? You only mention that SFT performs worse, but it’s not very clear how this is shown in the figure. For example, it seems that PG and Q-learning tend to hallucinate more than SFT — is that correct? I’m referring to the strong diagonal from bottom-left to top-right, which appears high for PG and Q-learning but low for SFT (as in the dataset). Also, Q-learning seems to produce a large bright block on the right — what does that correspond to?\n\n**Lack of empirical validation for important assumption**  \nTheoretical analyses throughout the paper hinge on Assumption 3.1 — motivated by Wang et al. (2024b) — that the model’s logits depend only on the current and target nodes. However, this assumption is adopted without further empirical verification. The paper does not test whether real Transformer models indeed satisfy this reduced dependency, leaving open whether the derived results and takeaways accurately reflect practical LLM behavior.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TWijIwriVZ", "reviewer": "Reviewer_BBiB", "summary": "This paper provides a theoretical analysis of reinforcement learning (RL) methods, particularly policy gradient (PG) and Q-learning, for planning in large language models (LLMs). It builds on recent empirical works showing that RL enhances reasoning and planning beyond supervised fine-tuning (SFT), and offers a rigorous mathematical explanation of why and how that happens.\n\nThe authors adopt a graph-based abstraction of planning, where path-finding over a directed graph represents multi-step reasoning or decision-making (e.g., tool use, theorem proving, or task sequencing). They theoretically and empirically analyze three regimes 1) SFT: Learns co-occurrence-based correlations without transitive reasoning, leading to memorization rather than true planning; 2) Policy Gradient (PG): Outperforms SFT via exploration-based data augmentation, but suffers from diversity collapse (reduced output variety even after convergence). KL regularization mitigates this but can hurt accuracy; 3) Q-Learning: Addresses PG’s pitfalls by enabling off-policy learning and preserving output diversity at convergence—but risks reward hacking unless process rewards (stepwise feedback) are used instead of binary outcome rewards. The authors validate these behaviors using both synthetic Erdos–Rényi graphs and the Blocksworld planning benchmark", "questions": "- How do the diversity collapse dynamics extend to multi-step, natural language reasoning, where outputs are discrete tokens instead of graph edges?\n\n- Can Q-learning be feasibly integrated with large transformer-based LLMs at scale? Are there existing implementations (e.g., in VeRL or R1) that align with the off-policy formulation analyzed here?\n\n- Would adding entropy regularization or temperature annealing alter Theorem 4.3’s diversity-collapse dynamics?\n\n- How sensitive are the conclusions to the assumption of persistent exploration? Could adaptive exploration decay or real-world sampling noise violate this condition?\n\n- Could you relate process rewards to verifiable subgoal rewards used in reasoning tasks (e.g., partial correctness of intermediate steps)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper has a strong theoretical foundation: The paper provides a clean, formal framework for analyzing planning in LLMs, translating abstract reasoning tasks into tractable graph-based path-planning problems. \n\n- Novel Theoretical Insights into RL for LLMs: Identifies diversity collapse in PG as a fundamental phenomenon, theoretically proving that entropy monotonically decreases after convergence—a widely observed empirical issue in RLHF and GRPO models. It reveals that Q-learning avoids this trap by updating state-action values rather than direct log-probabilities, preserving diversity across valid next actions. As well, it formalizes the reward hacking issue and demonstrates how process rewards eliminate it.\n\n- Empirical Validation with Clear Visual Evidence: Figure 1 visualizes adjacency reconstruction in Blocksworld—showing how SFT fails to recover the graph structure while PG and Q-learning succeed. Then, the paper connects theory to practice via Figures 2–4. Figure 2 shows how KL regularization balances accuracy and diversity in PG. Figure 3 demonstrates Q-learning’s superior test accuracy and stability. Figure 4 draws heatmaps of logits reveal that feasible transitions converge to uniform values (confirming theoretical diversity preservation).\n\n- Clear Conceptual Takeaways: The paper summarizes each section with concise takeaway statements, synthesizing intuition for both theoretical and empirical results. This makes it approachable despite heavy mathematical content.", "weaknesses": "- Limited Practical Scope: The paper primarily focuses on simplified models (one-layer, single-head transformers and synthetic graphs). While conceptually illuminating, the connection to real-world LLM training (e.g., PPO/GRPO pipelines) remains somewhat indirect. It has no experiments on large-scale reasoning benchmarks (e.g., MATH, GSM8K, or ALPINE) are included to show external validity.\n\n- Simplistic Reward Modeling: The outcome vs. process reward distinction is theoretically elegant but lacks discussion on how these correspond to real LLM reward signals (e.g., chain-of-thought correctness, factual consistency). Furthermore, reward hacking is discussed in the graph setting but not demonstrated empirically in language tasks.\n\n- Empirical Validation is Narrow: All empirical results use synthetic environments. If authors can have results on real text-based reasoning or tool-use tasks, then the statements would be much more strengthen that the same dynamics hold for modern RL-tuned LLMs.\n\n- Absence of Theoretical Analysis for PPO/GRPO: Although Section F shows PPO is mathematically equivalent to unclipped PG under stop-gradient treatment, this misses important effects like clipping or value normalization used in real systems.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)