Trajectory Bellman Residual Minimization: A Simple Value-Based Method for LLM Reasoning

Thank you for the review! Please see responses to individual questions below.

Q: Relation between TBRM and Policy Gradient

In comparison with policy gradient, the gradient of the TBRM loss can be decomposed as follows:

Where $\mu$ is the policy used to generate the trajectories, and $\overline{\theta}$ denotes a copy of $\theta$ whose parameters are held fixed (i.e., not updated). In the on-policy case, i.e., $\mu = \pi_\theta$, the first term in TBRM resembles the REINFORCE gradient but with a baseline $V_\theta(s_1)$ for variance reduction, and the second term corresponds to the gradient of the value function.

Q: Comparison between the Performance of TBRM and GRPO

We acknowledge that the performance of TBRM is comparable to GRPO. The main advantage of TBRM is that: it is simpler to implement, requiring neither importance-sampling ratios nor clipping; it is more flexible, as it naturally supports off-policy training and works effectively with a single rollout; and it provides theoretical guarantees, including finite-sample bounds. These strengths make TBRM a practical and principled alternative across a range of scenarios, maybe particularly in tasks involving multi-turn conversations and agentic workflows, where trajectory collection is both costly and challenging. We view adapting TBRM to such settings as a promising direction for future work.

In the next revision of our paper, we will revise the wording that suggests TBRM consistently outperforms GRPO and instead emphasize that the two methods achieve comparable performance. This clarification is intended to prevent potential confusion for readers.

Q: Ablation Study of State-Action-Level BRM

We agree with the reviewer on the importance of including an ablation study compared against the original BRM. To highlight the advantages of the trajectory-level approach, we implemented a state-action-level variant of BRM and ran experiments using the same hyperparameters as in our main setup on Qwen2.5-Math-1.5B. The results show that training reward quickly collapses, and the model outputs become random and meaningless. Intuitively, this degradation occurs because BRM has to propagate the sparse reward signal, which only receives at the final token, back through multiple token-wise regressions, whereas the TBRM provides a better implicit credit assignment through a single trajectory-level objective.

Q: Experiment Results with More Rollouts

We rerun GRPO and TBRM using most hyperparameters from DAPO [1] to ensure a fair comparison. Specifically, we used a prompt batch size of 512 and generated n=16 responses per prompt. For GRPO, we set the microbatch size to 512, resulting in 16 updates per training step. The experiments were conducted on the Qwen2.5-Math-7B model, following the same evaluation pipeline described in our paper. Both algorithms were trained for 100 steps. Each run takes approximately 40 hours on 4x H100 GPUs. Our results show that TBRM remains comparable to GRPO under these aligned settings.

[1] DAPO: An Open-Source LLM Reinforcement Learning System at Scale. Yu et al. 2025

Thank you for the review! Please see responses to individual questions below.

Q: Comparison with Other Related Algorithms

For a detailed comparison between TBRM and soft Q-learning, we refer the reviewer to the updated Appendix E of the supplementary material (please see Appendix E in the ZIP file of the supplementary material, which has been updated from the version in the main paper). The major differences include (1) TBRM employs direct optimization rather than an iterative approach, and (2) TBRM’s loss function operates on complete trajectories rather than summing losses over individual timesteps within trajectories, hence eliminating the need for per-step reward.

In comparison with policy gradient, the gradient of the TBRM loss can be decomposed as follows:

Where $\mu$ is the policy used to generate the trajectories, and $\overline{\theta}$ denotes a copy of $\theta$ whose parameters are held fixed (i.e., not updated). In the on-policy case, i.e., $\mu = \pi_\theta$, the first term in TBRM resembles the REINFORCE gradient but with a baseline $V_\theta(s_1)$ for variance reduction, and the second term corresponds to the gradient of the value function.

Finally, CoPG is basically optimizing the difference between the Bellman residuals of two trajectories. It falls under the category of “DAA-pair” algorithms discussed in Appendix E (please see Appendix E in the ZIP file of the supplementary material, which has been updated from the version in the main paper).

Q: Modeling LLM Rollout as MDP or Contextual Bandit

Our setting applies to more general deterministic MDPs than the token-level MDP, even when the transition dynamics are unknown a priori. This formulation naturally captures a wide range of scenarios beyond standard single-turn question answering. Notable examples include multi-turn conversations and agentic workflows [1], which cannot be adequately modeled as contextual bandits. For instance, outputs from external tools are part of the deterministic state transitions, so the transition is no longer merely appending new actions to the state. Our algorithm, with MDP setting, correctly avoids taking gradients on the output from external tools, whereas a naive contextual bandit approach would incorrectly take gradients through them. This limitation of contextual bandits has been empirically demonstrated in [2].

Additionally, while token-level MDPs are statistically equivalent to contextual bandits in some cases, contextual bandit algorithms may be computationally inefficient when applied to token-level settings, for example, when they require explicit optimization over the action space, which corresponds to all possible token sequences.

Q: Limitations of Single-Sample TBRM on Challenging Datasets

We agree that on sufficiently challenging datasets, TBRM may require multiple samples to achieve strong performance. However, our claim that "TBRM only requires one sample per prompt" refers to the flexibility of its objective formulation, rather than its statistical efficiency. We will clarify this distinction in the next revision to avoid confusion.

Q: Better Benchmarking with Hyperparameters from DAPO

We rerun GRPO and TBRM using most hyperparameters from DAPO [3] to ensure a fair comparison. Specifically, we used a prompt batch size of 512 and generated n=16 responses per prompt. For GRPO, we set the microbatch size to 512, resulting in 16 updates per training step. The experiments were conducted on the Qwen2.5-Math-7B model, following the same evaluation pipeline described in our paper. Both algorithms were trained for 100 steps. Each run takes approximately 40 hours on 4x H100 GPUs. Our results show that TBRM remains comparable to GRPO under these aligned settings.

Q: Clarification on logits used in the algorithm

$\mathsf{logit}_\theta (s_1,a_1)$ is the raw logit (before softmax) of the first output token with input being a prompt from the dataset.

We agree with the reviewer that providing a more precise description of the algorithm and a clearer definition of the Q-function would improve readability. We will revise the wording accordingly in our next revision.

[1] Search-r1: Training llms to reason and leverage search engines with reinforcement learning. Jin et al. 2025

[2] Building math agents with multi-turn iterative preference learning. Wei et al. 2024

[3] DAPO: An Open-Source LLM Reinforcement Learning System at Scale. Yu et al. 2025

Thank you for the review! Please see responses to individual questions below.

Q: Comparison between TBRM and GRPO with $n=4$

We acknowledge that TBRM and GRPO achieve comparable performance on Qwen2.5-Math-7B when using 4 rollouts. However, TBRM offers several notable advantages: it is simpler to implement, avoiding the need for importance sampling ratios or clipping; it is more flexible, naturally accommodating off-policy training and performing effectively even with a single rollout; and it is theoretically grounded, providing finite-sample guarantees. These properties make TBRM a practical and principled choice across a wide range of settings.

To prevent potential confusion, we will clarify in the introduction that TBRM achieves its best performance when using more than one rollout ($n > 1$).

Q: Extension to Non-Mathematical Tasks

To demonstrate the generalizability of our method beyond mathematical tasks, we evaluate it on five tasks from the reasoning-gym [1] under the "graphs" category: course_schedule, family_relationships, largest_island, quantum_lock, and shortest_path. These tasks are naturally represented as graphs, consisting of nodes and edges, and typically require traversing connections to identify relationships, compute optimal paths, or determine reachable components. They involve reasoning patterns that differ significantly from those in mathematical tasks.

We construct a training set of 10,000 problems, with 2,000 questions per task, and a test set of 500 problems, comprising 100 questions from each task. For both training and evaluation, we use the official verifiers provided by reasoning-gym to compute rewards. Our experiments are conducted on Qwen2.5-Math-1.5B using both TBRM and GRPO, with a prompt batch size of 1024 and 4 sampled responses per question (n = 4). Models are trained for 100 steps. All evaluations are conducted using greedy decoding. Results demonstrate that TBRM generalizes well to diverse reasoning tasks and performs on par with GRPO.

Q: Clarification on Training Set DAPO

We thank the reviewer for raising this important point. Our training data is an open-source dataset (DAPO [2]), which is widely used in the community. We did not curate or modify this training set, and our work does not aim to contribute to dataset construction. Instead, our focus is on the methodological advancement evaluated on standard benchmarks.

To mitigate concerns about potential train-test leakage, we selected well-established benchmarks as our test sets, which are publicly available and commonly used for evaluation. However, since we rely on existing open-source data for training, we acknowledge that we cannot guarantee the absence of any incidental overlap (e.g., paraphrased or similar examples) without access to the full provenance of the dataset.

[1] REASONING GYM: Reasoning Environments for Reinforcement Learning with Verifiable Rewards. Stojanovski et al. 2025

[2] DAPO: An Open-Source LLM Reinforcement Learning System at Scale. Yu et al. 2025

Thank you for the review! Please see responses to individual questions below.

Q: Limitations of Deterministic MDPs in Algorithmic Proofs

Our algorithm is designed for discrete, tokenized settings relevant to language models. While the guarantee for Algorithm 1 (Theorem 2) only works for deterministic MDPs, the change-of-trajectory-measure lemma (Lemma 1) holds for stochastic MDPs as well. And this lemma has broader use beyond the analysis of the algorithm in this paper. Extending our algorithm and its analysis to the general MDP setting is beyond the scope of this work and is left for future research.

Q: Missing Error Bars/Confidence Interval

We rerun the evaluation five times with different random seeds and compute confidence intervals following the methodology used in arena-hard-auto [1]. We report the 90% confidence interval, using the 0.05 and 0.95 quantiles as the lower and upper bounds, respectively.

Q: Relation between TBRM Loss and L2 Loss of AWR

In our Eq.3, logits are effectively Q-function, and $\pi_\theta$ are effectively the softmax policy of logits/Q-function. We will then explain the difference from the $L_2$ version of AWR based on this.

If our understanding is correct, the $L_2$ version of AWR is doing Monte Carlo style regression, regressing the on-policy log-probability of soft-max policy of return-to-go, namely $\log\mathsf{softmax}\circ R_{\rm rtg}(s ,a)=\log \frac{\exp\{R_{\rm rtg}(s ,a)\}}{Z(s)}$, to $\log\pi(a \mid s)$, where $Z(s)$ is the partition function. In this way, one must do the regression per state-action, which requires to handle the normalization factor $Z(s)$ per state-action.

On the other hand, TBRM is closer to the TD-style algorithm, which is naturally off-policy. TBRM loss is in the trajectory level, and the $\log\pi_\theta$ terms in the middle plays the same role of bootstrap in as in typical TD approaches, which fixes the off-policy data.

We suspect that the reviewer might also question if $L_2$ version of AWR in KL-regularized RL would close that gap. It still cannot, because in that case 1) $L_2$ version of AWR should use $\log\pi_{t-1}$ (log-prob of the policy that collected the data) over return-to-go, i.e., $R_{\rm rtg}(s_h,a_h)=\sum_{h'=h}^H \left( r(s_{h'},a_{h'})-\beta \log\frac{\pi_{t-1}(a_{h'}\mid s_{h'})}{\pi_{\mathsf{ref}}(a_{h'}\mid s_{h'})}\right)$, but TBRM should use $\log\pi_\theta$ (log-prob of current policy), i.e., $R_{\rm rtg}(s_h,a_h)=\sum_{h'=h}^H \left( r(s_{h'},a_{h'})-\beta \log\frac{\pi_\theta(a_{h'}\mid s_{h'})}{\pi_{\mathsf{ref}}(a_{h'}\mid s_{h'})}\right)$. They are same in purely on-policy but not for off-policy; 2) $L_2$ version of AWR should still have squared loss per-state-action, but TBRM has trajectory-level loss.

Q: Extension to Non-Mathematical Tasks

To demonstrate the generalizability of our method beyond mathematical tasks, we evaluate it on five tasks from the reasoning-gym [2] under the "graphs" category: course_schedule, family_relationships, largest_island, quantum_lock, and shortest_path. These tasks are naturally represented as graphs, consisting of nodes and edges, and typically require traversing connections to identify relationships, compute optimal paths, or determine reachable components. They involve reasoning patterns that differ significantly from those in mathematical tasks.

We construct a training set of 10,000 problems, with 2,000 questions per task, and a test set of 500 problems, comprising 100 questions from each task. For both training and evaluation, we use the official verifiers provided by reasoning-gym to compute rewards. Our experiments are conducted on Qwen2.5-Math-1.5B using both TBRM and GRPO, with a prompt batch size of 1024 and 4 sampled responses per question (n = 4). Models are trained for 100 steps. All evaluations are conducted using greedy decoding. Results demonstrate that TBRM generalizes well to diverse reasoning tasks and performs on par with GRPO.

[1] From Crowdsourced Data to High-Quality Benchmarks: Arena-Hard and BenchBuilder Pipeline. Li et al. 2024

[2] REASONING GYM: Reasoning Environments for Reinforcement Learning with Verifiable Rewards. Stojanovski et al. 2025