Step Rejection Fine-Tuning: A Practical Distillation Recipe Igor Slinko, Ilia Zavidnyi, Egor Bogomolov, Yaroslav Zharov
arXiv:2605.10674v1 [cs.LG] 11 May 2026
JetBrains Research • igor.slinko, ilia.zavidnyi, egor.bogomolov, [email protected]
Rejection sampling Fine-Tuning (RFT) is a standard method for training LLM agents, where unsuccessful trajectories are discarded from the training set. In the context of SWE-bench tasks, this corresponds to filtering out runs where the submitted patch does not pass the tests. However, this approach discards unresolved trajectories, even though they form a large portion of all trajectories for hard tasks and even then may be partially correct. In this work, we propose Step Rejection Fine-Tuning (SRFT)—a practical way to leverage these unresolved trajectories. For this, we employ a critic LLM to assess the correctness of each step in a trajectory. Consequently, during training, we mask the loss for erroneous steps while retaining them in the context window. This way we ensure the model learns to recover from errors without reproducing them. Evaluation on SWE-bench Verified shows that while RFT improves the resolution rate by 2.4% by excluding unresolved trajectories, SRFT improves it by 3.7% by filtering them instead of discarding completely, reaching the total resolution rate of 32.2%. Date: April 2026
1 Introduction LLM-based agents are systems designed to autonomously perceive and interact with environments to achieve complex goals (Xi et al., 2023; Wang et al., 2023), with the ability to reason, plan, and use tools to solve open-ended problems. LLM-based agents often employ paradigms like ReAct (Yao et al., 2023) or Reflexion (Shinn et al., 2023) that interleave reasoning and acting phases, constituting the so-called trajectory that consists of multiple such steps. Rejection sampling Fine-Tuning (RFT) (Yuan et al., 2023) has emerged as a standard paradigm for training LLM agents (Pan et al., 2025; Jain et al., 2025; Yang, Lieret, et al., 2025). This process involves generating multiple trajectories for a given task, filtering for those that successfully resolve the task, and performing supervised fine-tuning on this filtered set. Filtering is needed, since incorporating unresolved trajectories into the training set degrades performance by teaching the model to imitate erroneous actions. A major limitation of RFT is the inability to leverage unresolved trajectories. For instance, SWE-smith (Yang, Lieret, et al., 2025) introduces a large-scale dataset of agent trajectories obtained by solving synthesized software engineering tasks and a model trained on this dataset. However, due to RFT, they discard approximately 61% of the collected runs, therefore losing a lot of potentially informative data. Our key insight is that unresolved trajectories are not entirely erroneous; rather, they often consist of correct and useful steps interspersed with errors. Our manual analysis of 20 trajectories indicates that even in unresolved trajectories, only up to 24% of steps can be classified as mistakes. To capitalize on this finding, we introduce Step Rejection Fine-Tuning (SRFT). In this method, we utilize a critic Large Language Model (LLM) to discriminate trajectories on a lower level, marking singular steps of trajectories as either worthy or unworthy to train on. This method allows the model to learn from the valid portions of unresolved trajectories without internalizing mistakes. The contributions of this paper are as follows. We present a practical light-weight approach that allows utilizing the data routinely considered noise without changing the overall pipeline. The approach is detailed in Section 3. Then, we conduct an experiment that shows that on the challenging SWE-bench Verified benchmark, our method successfully improves over the naïve RFT baseline. The experiments 1
are described and analyzed in Section 4.
2 Related Work Research has increasingly moved beyond standard Rejection sampling Fine-Tuning—which treats all actions in a resolved trajectory as equally valid—towards extracting granular supervision from suboptimal or failed attempts. Learn-by-interact (Su et al., 2025), an instance of a broader Hindsight Experience Replay (HER) approach (Andrychowicz et al., 2017), addresses instruction–trajectory misalignment in long trajectories. They propose a backward construction method that decomposes trajectories into shorter segments and synthesizes specific tasks for each, thereby creating valid demonstrations for these synthesized tasks. While this and other HER-based methods effectively utilize available data, they rely on generating synthetic instructions. In contrast, our approach focuses on filtering steps within the original task context, avoiding the need for task synthesis. Also, our method alleviates the risk of drifting off of the ground truth task. SWEET-RL (Zhou et al., 2025) tackles the credit assignment problem by transforming trajectorylevel feedback into step-wise signals. It trains a step-wise critic by comparing pairs of trajectories and constraining the score to be a sum of per-step contributions, forcing an implicit decomposition of preferences. This critic then guides policy optimization via Direct Preference Optimization (DPO; Rafailov et al., 2023). While effective, this introduces the complexity of Reinforcement Learning. In contrast, our method avoids training a critic with complex objectives, offering a simpler alternative within the supervised learning paradigm. Finally, STeP (Chen et al., 2025) employs a teacher-in-the-loop to actively synthesize “self-reflected” trajectories where errors are immediately followed by teacher-generated reflections and corrections. By applying partial masking to error steps, they enable the model to learn recovery strategies. While our work shares the core mechanism of partial masking, we propose a simple practical approach that can retroactively fit into existing pipelines. Specifically, instead of requiring expensive real-time teacher intervention to synthesize new correction steps, we use an offline critic to salvage valid signals from standard unresolved trajectories.
3 Method 3.1 Preliminaries We consider the problem of distilling the agentic behavior from a strong Teacher model to a weaker Student model. The distillation process consists of gathering the outputs of the Teacher and training the Student to mimic them. We denote the set of the collected Teacher outputs as D. We denote the Student model as a function fθ parametrized by weights θ, that estimates the probability of the next token given the previous tokens. For agentic behavior in particular, we need to sample coherent ReAct trajectories rather than oneshot predictions, which complicates the dataset collection. Each trajectory τi ∈ D is defined as τi = (s, ui , (ai,0 , oi,0 ), (ai,1 , oi,1 ), . . . , (ai,Ti , oi,Ti )), comprising a system message s, a task description ui , and a i sequence of assistant actions and corresponding environment observations {(at , ot )}Tt=0 , where Ti is the length of the i-th trajectory in steps. For the sake of simplicity, we use τi,[0:t] to denote the concatenation of the trajectory τi up to but not including the t-th step, and at,[0:j] to denote the string representation of the action at up to but not including j-th token.
3.2 Methods Investigated in Current Work The naïve distillation is performed by minimizing the following loss L(θ, D), where NLL is the negative log-likelihood loss. |at | Ti X XX L(θ, D) = NLL fθ (τi,[0:t] + at,[0:j] ), at,j τi ∈D t=0 j=0
2
The RFT filters the raw teacher trajectories in D to retain only trajectories satisfying the final success criteria of the environment, resulting in the subset Ds . For example, in SWE-bench, it means that after the agent has finished, a pre-defined set of tests should succeed. The Student model is then trained by minimizing the loss L(θ, Ds ). However, this approach discards the set of unresolved trajectories Df = D \ Ds , which may be large for complex tasks. To leverage the data rejected under the RFT framework, we introduce a fine-grained supervision mechanism SRFT. Instead of discarding or keeping the entire trajectories, we introduce weights W to keep alongside the trajectories. Weight wi,t ∈ W corresponds to the step (ai,t , oi,t ) in trajectory τi . Given the weights, we modify the loss to perform the weighted distillation on the whole dataset with loss L(θ, D, W). LW (θ, D, W) =
Ti X X
wi,t ·
τi ∈D ′ t=0
|at | X
NLL fθ (τi,[0:t] + at,[0:j] ), at,j
j=0
3.3 SRFT Instantiation in Current Work SRFT supports different weighting schemes, but for our current experiments, we instantiate it as follows. For successful trajectories in Ds , we assign a weight of 1 to each step. For unresolved trajectories in Df , we employ a critic model to label each action at . We classify a step as good if it advances the task, harmful if it hinders progress (e.g., introduces a bug), and unnecessary if it neither helps nor damages. We group good and unnecessary steps into a single category of Productive steps. Accordingly, we assign wt = 0 to harmful steps and wt = 1 otherwise, forming Wf . The total set of weights is W = Wf ∪ Ws . Given this weighting scheme, the loss LW simply omits the loss calculation for the steps that were marked as harmful. Conveniently, for binary weights, this corresponds to modifying token-wise masks. This is routinely done for observations, and including masks for some pre-defined steps incurs little engineering overhead. A further example of how this is applied to a trajectory is depicted in Figure 1. L
RFT:
s
u
a0
L
o0
a1
L
SRFT:
s
u
a0
L
a2
o1
L
o0
a1
Context Flow
L
o2
···
aT -1
a2
Loss Contribution
oT -1
L
L
o1
L
o2
···
aT -1
w=1
aT
L
oT -1
aT
w=0
Figure 1. Comparison of RFT and SRFT training approaches.
4 Experimental Setup To experimentally demonstrate that SRFT improves the performance of the Student model, we employ it on a complex Software Engineering task. We fine-tune Qwen2.5-Coder-32B-Instruct (Hui et al., 2024) as our Student model, using the SWE-Agent (Yang, Jimenez, et al., 2024) framework as the scaffold. For training, we utilize the SWE-smith-trajectories dataset (Yang, Lieret, et al., 2025), which consists of approximately 25,000 trajectories generated by SWE-agent. The dataset comprises 39% resolved and 61% unresolved trajectories. For our study, we sample a balanced set of 5,000 resolved (Ds ) and 5,000 unresolved (Df ) trajectories. We further instantiate datasets for different methods as described in Section 3. To provide the markup for the SRFT method, we employ Claude 4 Sonnet (snapshot 20250514) (Anthropic, 2025) as the critic model, incurring a total cost of $660 for 5,000 trajectories. The critic prompt was selected by expert labeling of trajectories, and is presented alongside the expert data in Section A. Table 1 details the distribution of step labels. We note that unresolved trajectories contain more steps marked as 3
Label
Resolved
Unresolved
Productive steps Harmful steps
95.9% 4.1%
93.0% 7.1%
Table 1. Label distribution for 5,000 resolved and 5,000 unresolved trajectories. Training Data
Resolved (%)
Base model Naïve distillation RFT SRFT
7.0 ± 1.3 (-21.5) 28.5 ± 1.7 30.9 ± 1.1 (+2.4) 32.2 ± 0.9 (+3.7)
Table 2. Main results on SWE-bench Verified. Each experiment was run 7 times; we report mean ± standard deviation. Base model performance is from the paper R2E-Gym (Jain et al., 2025), which uses the same scaffold and model.
harmful, but the critic was never provided with the resolution status of the trajectory. This aligns with the intuition that the unresolved trajectories contain potentially adversarial patterns not to be distilled. We evaluate the performance of the Student model on the SWE-bench Verified dataset (OpenAI, 2024). To ensure robustness, each experiment is repeated 7 times. We report the Resolved Rate in percent. The results are presented in Table 2. Consistent with our hypothesis, naively including unresolved trajectories leads to a performance degradation compared to the RFT baseline (28.5% vs 30.9%), as the model internalizes errors present in the failed attempts. However, by applying critic-guided masking, we not only mitigate this degradation but achieve a performance gain, outperforming the RFT baseline (32.2% vs 30.9%). This improvement is statistically significant; a bootstrap analysis confirms a gain of 1.3% with a 95% confidence interval of [0.4, 2.3] (refer to Section B for detailed statistical analysis).
5 Limitations and Future Work Our approach relies on the accuracy of the critic. Mislabeling valid steps as harmful can reduce the effective training data, while failing to identify subtle errors can allow them to propagate into the student model. We leave a thorough study of different critics and labeling approaches to future work. In this work, we did not explore the possibility of the loss weights beyond binary. However, we note that this may be the key to further improving the quality, both on resolved and unresolved trajectories. While we test SRFT on a challenging SWE-bench task, we acknowledge the lack of a generalization study and leave the investigation of the method’s generalizability to future work.
6 Conclusion We have presented Step Rejection Fine-Tuning, a straightforward yet effective enhancement to the Rejection sampling Fine-Tuning distillation method that unlocks the value of unresolved trajectories. By selectively masking steps, we enable agents to learn from the partial successes of the Teacher model within failed attempts without internalizing its errors. On SWE-bench Verified, our method yields statistically significant improvements over the RFT (32.2% vs 30.9% resolved issues), highlighting the potential of step-level supervision in agentic distillation.
4
References Andrychowicz, Marcin et al. (2017). “Hindsight Experience Replay”. In: Advances in Neural Information Processing Systems. Ed. by I. Guyon et al. Vol. 30. Curran Associates, Inc. url: https://proceedings. neurips.cc/paper_files/paper/2017/file/453fadbd8a1a3af50a9df4df899537b5-Paper.pdf. Anthropic (2025). Claude 4 System Card. Tech. rep. Model version: claude-4-sonnet-20250514. Anthropic. url: https://www.anthropic.com/claude-4-system-card. Chen, Yihan et al. (2025). “Training LLM-Based Agents with Synthetic Self-Reflected Trajectories and Partial Masking”. In: arXiv preprint arXiv:2505.20023. url: https://arxiv.org/abs/2505.20023. Hui, Binyuan et al. (2024). “Qwen2.5-Coder Technical Report”. In: arXiv preprint arXiv:2409.12186. url: https://arxiv.org/abs/2409.12186. Jain, Naman et al. (2025). “R2E-Gym: Procedural Environments and Hybrid Verifiers for Scaling OpenWeights SWE Agents”. In: arXiv preprint arXiv:2504.07164. url: https://arxiv.org/abs/2504. 07164. OpenAI (2024). SWE-bench Verified. Web page. Accessed: 2025-08-11. url: https://openai.com/index/ introducing-swe-bench-verified/. Pan, Jiayi et al. (2025). “Training Software Engineering Agents and Verifiers with SWE-Gym”. In: Proceedings of the 42nd International Conference on Machine Learning (ICML 2025). arXiv:2412.21139, accepted at ICML 2025. url: https://arxiv.org/abs/2412.21139. Rafailov, Rafael et al. (2023). “Direct Preference Optimization: Your Language Model is Secretly a Reward Model”. In: arXiv preprint arXiv:2305.18290. url: https://arxiv.org/abs/2305.18290. Shinn, Noah et al. (Mar. 2023). “Reflexion: Language Agents with Verbal Reinforcement Learning”. In: arXiv e-prints. doi: 10.48550/arXiv.2303.11366. url: https://arxiv.org/abs/2303.11366. Su, Hongjin et al. (2025). “Learn-by-interact: A Data-Centric Framework for Self-Adaptive Agents in Realistic Environments”. In: arXiv preprint arXiv:2501.10893. url: https://arxiv.org/abs/2501. 10893. Wang, Lei et al. (Aug. 2023). “A Survey on Large Language Model based Autonomous Agents”. In: arXiv e-prints. doi: 10.48550/arXiv.2308.11432. url: https://arxiv.org/abs/2308.11432. Xi, Zhiheng et al. (2023). “The Rise and Potential of Large Language Model Based Agents: A Survey”. In: CoRR abs/2309.07864. doi: 10.48550/arXiv.2309.07864. url: http://arxiv.org/abs/2309.07864. Yang, John, Carlos E Jimenez, et al. (2024). “SWE-agent: Agent-Computer Interfaces Enable Automated Software Engineering”. In: arXiv preprint arXiv:2405.15793. url: https://arxiv.org/abs/2405. 15793. Yang, John, Kilian Lieret, et al. (2025). SWE-smith: Scaling Data for Software Engineering Agents. arXiv: 2504.21798 [cs.SE]. url: https://arxiv.org/abs/2504.21798. Yao, Shunyu et al. (2023). “ReAct: Synergizing Reasoning and Acting in Language Models”. In: International Conference on Learning Representations. url: https://openreview.net/forum?id=WE_vluYULX. Yuan, Zheng et al. (2023). “Scaling Relationship on Learning Mathematical Reasoning with Large Language Models”. In: arXiv preprint arXiv:2308.01825. url: https://arxiv.org/abs/2308.01825. Zhou, Yifei et al. (2025). “SWEET-RL: Training Multi-Turn LLM Agents on Collaborative Reasoning Tasks”. In: arXiv preprint arXiv:2503.15478. url: https://arxiv.org/abs/2503.15478.
5
A Critic Prompt and Evaluation We validated the critic prompt by manually labeling a ground truth set of 20 trajectories (comprising 444 steps in total). Table 3 details the performance of the critic against these manual labels. The critic demonstrates robust performance in identifying good steps (F1=0.737), which constitutes the majority class. While performance on minority classes (harmful and unnecessary) is lower, the overall system benefit suggests that even imperfect filtering is preferable to naïve data inclusion or complete data exclusion. Label
F1
Precision
Recall
good unnecessary harmful
0.737 0.236 0.384
0.660 0.253 0.604
0.835 0.221 0.282
Accuracy
0.586
Table 3. Critic prompt evaluation on 444 manually labeled steps.
The specific prompt used for the critic is provided below: You are evaluating SWE-agent trajectories. Your goal is to annotate each step of the trajectory. This annotation will be used to train a model only on good steps. First, construct a brief "ideal solution path" and patch (in your own reasoning, not in the output). Then, for each step, decide if it matches or deviates from this ideal path. ## Mark Each Step - **"good"**: Part of or necessary for the ideal shortest path (exploration, correct fixes). - **"harmful"**: Moves away from the path (wrong edits, irrelevant files), harmful tool usage, etc. - **"unnecessary"**: Not harmful but not part of the path (redundant exploration, repeated actions). ## Rules - Always compare each step to the shortest correct path you imagined. - A step is "good" only if it is part of or necessary for that shortest path. - A step is "harmful" if it deviates from or worsens the trajectory relative to that path. - A step is "unnecessary" if it neither helps nor harms but is outside the shortest path.
6
B Statistical Significance and Extended Results Given the inherent high variance in LLM agent evaluation, we conducted a rigorous statistical analysis. We identified two primary sources of noise: rollout variance, where two distinct rollouts of the same model can differ by up to 4.8%, and training variance, where two models trained with different seeds can differ by 1.1% on average across 7 rollouts. To verify that our method provides a statistically significant improvement, we focused on the unresolved split, which most clearly highlights the difference between using all steps versus masking. We performed 5 training runs: 3 without masking and 2 with masking. Accounting for both sources of noise, we still observe a performance increase (see Table 4). As shown in Figure 2, the bootstrap analysis confirms a statistically significant improvement of 1.1% with a 95% confidence interval of [0.4, 1.8].
Figure 2. Bootstrap analysis of Resolved Rate% for 5k Unresolved vs 5k Unresolved masked. Configuration
Resolved%
pass@7%
5k Unresolved (train #1) 5k Unresolved (train #2) 5k Unresolved (train #3) 5k Unresolved (combined)
28.8 ± 1.0 27.7 ± 0.9 29.1 ± 1.0 28.5 ± 1.1
42.8 42.2 41.8 42.3
5k Unresolved masked (train #1) 5k Unresolved masked (train #2) 5k Unresolved masked (combined)
29.8 ± 0.6 29.5 ± 1.5 29.7 ± 1.1
42.6 43.8 43.2
Table 4. Results from multiple training runs on 5k Unresolved trajectories. Mean and standard deviation are calculated over 7 rollouts for individual runs; combined results are calculated over 21 and 14 rollouts respectively. Training Data
Resolved%
pass@7%
7.0 ± 1.3
Base model RFT Naïve distillation SRFT
5k Resolved 5k Unresolved 5k Resolved + 5k Unresolved 5k Resolved + 5k Unresolved (masked) 5k Resolved (masked) + 5k Unresolved (masked)
30.9 ± 1.1 27.7 ± 0.9 28.5 ± 1.7 32.2 ± 0.9 29.4 ± 1.6
Table 5. Extended experimental results on SWE-bench Verified.
7
45.8 42.2 43.8 45.8 41.2