Source-linked AI summary
Scaling LLM Multi-turn RL with End-to-end Summarization-based Context Management
Miao Lu, Weiwei Sun, Weihua Du, Zhan Ling, Xuesong Yao, Kang Liu, Jiecao Chen
TL;DR
Long-horizon multi-turn RL for LLM agents is constrained by rapidly growing tool-use context. The paper introduces end-to-end summarization-based context management and SUPO, which improves task performance while maintaining compact context and supports scaling summarization beyond training.
Problem
Long-horizon LLM-agent RL accumulates prompt, output, tool-observation, and reasoning context, creating a fundamental training bottleneck.
Method
The paper periodically replaces tool-use history with task-relevant LLM-generated summaries and optimizes summarization jointly with tool-use behavior through a policy-gradient formulation, instantiated as SUPO.
Results
SUPO achieves higher CodeGym performance than GRPO at the same effective context length, reaches 53% on BrowseComp-Plus, and reaches 60.0% when scaling summary rounds beyond training.
Takeaways & Limitations
Summarization-based context management provides a scalable way to train long-horizon RL agents beyond fixed context limits while preserving task-relevant rollout information.
Takeaways & Limitations
The framework specifically incorporates summarization-based context management into the decision pipeline; obtaining correct policy gradients for more general agentic workflows remains future work.
Abstract
from arXiv · showhide
We study reinforcement learning (RL) fine-tuning of large language model (LLM) agents for long-horizon multi-turn tool use, where context length quickly becomes a fundamental bottleneck. Existing RL pipelines can suffer from degraded instruction following, excessive rollout costs, and most importantly, strict context limits. To address these challenges, we introduce summarization-based context management to training. In specific, it periodically compresses the tool using history by LLM-generated summaries that retain task-relevant information to keep a compact context while enabling the agent to scale beyond the fixed context window. Building on this formulation, we derive a policy gradient representation that seamlessly enables standard LLM RL infrastructures to optimize both tool-use behaviors as well as summarization strategies in an end-to-end fashion. We instantiate this framework with \underline{SU}mmarization augmented \underline{P}olicy \underline{O}ptimization (\texttt{SUPO}), an LLM RL algorithm that enables long-horizon training beyond a fixed context limit. Experiments on interactive function calling and searching tasks demonstrate that \texttt{SUPO} significantly improves the success rate while maintaining the same or even lower working context length compared to baselines. We also demonstrate that for complex searching tasks, \texttt{SUPO} can further improve the evaluation performance when scaling test-time maximum round of summarization beyond that of training time. Our results establish summarization-based context management as a principled and scalable approach for training RL agents beyond a fixed context length limit.
1 Introduction
Long-horizon multi-turn RL is limited by rapidly growing context, which degrades instruction following, increases rollout costs, and imposes fixed context limits. The paper addresses this with end-to-end summarization-based context management and SUPO, which jointly optimize tool use and summarization.
- Motivation: Long-horizon tool-use RL accumulates context across prompts, outputs, observations, and reasoning traces, creating a scalability barrier.The resulting difficulties include degraded instruction following, excessive rollout costs, and fixed context limits.
- Approach: Summarization-based context management periodically compresses tool-use history into task-relevant LLM-generated summaries while keeping the working context compact.The context is reset to the initial prompt plus a summary, and the summary is optimized as part of the policy.
- Approach: The framework formalizes summarization as part of the MDP and derives a policy gradient that jointly optimizes tool-use behavior and summarization strategies.This formulation is intended to work with standard LLM RL infrastructure.
- Empirical validation: SUPO improves success rates on CodeGym and BrowseComp-Plus while using the same or shorter working contexts than baselines.The reported absolute improvements are +3.2% and +14.0% across the two environments.
2 Related Works
Prior RL methods for multi-turn tool use are largely bounded by the fixed context length of the fine-tuned LLM. This work extends that setting with end-to-end RL-trained summarization-based context management, related to but distinct from other learned memory approaches.
- RL for LLM tool use: Existing RL approaches for multi-turn tool use generally remain limited by the fixed context length of the underlying LLM.That bound limits the difficulty of tasks they can solve.
- Context management: This work augments the original RL modeling with end-to-end summarization-based context management to enlarge training beyond the model’s context limit.The approach belongs to working-context compression through LLM summarization.
- RL-based memory: Related RL memory methods learn summary or memory operations, whereas this work focuses on integrating summarization into multi-turn RL training for tool-use agents.The paper positions MemAgent, MEM1, and Memory-R1 as closely related efforts.
3 Preliminaries
The paper models multi-turn tool use as an MDP and extends it with summarization-based state transitions that periodically compress accumulated context. It derives a decomposed policy gradient so standard RL infrastructure can optimize the resulting trajectories under bounded working context.
- Standard MDP: The standard MDP represents accumulated tokens as states, LLM outputs as autoregressive actions, tool executions as observations, and final-answer verification as the reward.The state grows by concatenating each action and tool observation, while the reward is task-specific and rule-based.
- Summarization-based MDP: With a good summary strategy, the compressed context can preserve information needed to solve tasks whose required context exceeds the working limit.The framework jointly improves task completion through reasoning and tool calling and summarization for the task.
- Summarization-based MDP: Summarization-based context management triggers an LLM summary when accumulated context reaches a threshold, then restarts the working context with the initial prompt and summary.The transition rule distinguishes ordinary tool-use continuation from summary-triggered compression.
- Context bound: The working context is bounded by L + 2L_A + L_O + |v_sum|, combining the summarization threshold, LLM-call output allowance, tool-observation allowance, and summary-prompt length.This gives a quantitative characterization of context usage under the summarization process.
- Policy gradient: The policy-gradient theorem splits a rollout into complete trajectories whose gradients are summed, allowing each segment to be handled by existing vanilla multi-turn RL infrastructure.Each segment uses the initial prompt plus the previous trajectory’s summary as its new prompt.
- Scope: The presented framework is a special case of policy gradients for more general agentic workflows, leaving correct gradients for richer context management and multi-agency pipelines as future work.The stated goal is to demonstrate context management as part of the decision pipeline for scaling effective training context.
4 End-to-end RL Training of Agent with Summarization
SUPO adapts policy optimization to summarization-managed rollouts, allowing standard RL infrastructure to train agents beyond fixed context limits. Its design treats summarized rollout segments as trajectories, uses shared group-relative advantages, masks overlong failures, and controls context length near the summarization threshold.
- Overall Algorithm: SUPO: SUPO extends GRPO-style policy optimization with summarization-based context management to scale RL training beyond the LLM working context limit.The objective optimizes the LLM policy under a behavior policy while integrating summarization into rollout management.
- Trajectory management: Each rollout is split at summarization indices into multiple complete trajectories, allowing existing single-trajectory RL infrastructure to process the segments directly.Every split trajectory begins with the initial task prompt and the LLM-generated summary of preceding interactions.
- Advantage estimation: The same group-relative advantage is assigned to every token across a rollout’s split trajectories because all segments share the rollout reward.The estimator is computed relative to the rollout group, while an alternative trajectory-group estimator is evaluated separately.
- Overlong masking: Overlong masking excludes rollouts that fail to produce the final answer within the maximum step or summarization limits, avoiding pressure to suppress useful long-horizon summarization patterns.Without masking, the objective can be biased against long successful strategies and may cause summarization-pattern collapse.
- Fine control of context length: Discarding the action-observation pair that triggers summarization keeps the resulting working context within approximately the summary threshold plus summary length and action length.This makes the RL training context length approximately match the threshold when action sequences are shorter than the training context.
5 Experiments
Experiments on CodeGym and BrowseComp-Plus evaluate SUPO against GRPO and ablations, showing improved performance, longer tool-use behavior, and learned summarization strategies. Additional analyses examine overlong masking, advantage estimation, summary quality, and test-time scaling beyond training-time summarization rounds.
- Training and Evaluation Results of SUPO: SUPO achieves higher CodeGym scores than GRPO at the same 32K effective context, while reaching 53% on BrowseComp-Plus, 14% above GRPO.The comparison includes SUPO and ablations, with SUPO using 4K working context on CodeGym and 64K on BrowseComp-Plus.
- Further Analysis of SUPO: Without overlong masking, summarization collapses toward single-trajectory rollouts and conditional success falls to 0 during training on both tasks.These dynamics support the effectiveness of overlong masking in maintaining summarization-based long-horizon training.
- Further Analysis of SUPO: On BrowseComp-Plus, SUPO permits up to 3× more tool calling than GRPO during training, whereas GRPO’s average tool-calling count decreases.SUPO without overlong masking also shows a rapid drop in average tool calling.
- Further Analysis of SUPO: The advantage estimator used by SUPO outperforms the trajectory-group alternative, whose relative advantages for long successful trajectories are weakened by additional score-1 trajectories.The paper connects this intuition to worse test-time summarization-round scaling for the alternative estimator.
- Further Analysis of SUPO: The trained summarizer increasingly retains task-critical state, such as an exact array index in CodeGym and answer-bearing names from a searched article in BrowseComp-Plus.Before SUPO, the summaries omitted information needed to continue solving or answer the task.
- Scaling beyond Trajectory Number during Training: When test-time summarization rounds exceed training-time rounds, SUPO reaches the highest final accuracy among the compared algorithms, converging to 60.0%.This result is reported for BrowseComp-Plus and supports scaling summarization-based computation beyond the training configuration.
6 Conclusions and Future Works
The paper presents an RL framework that integrates summarization into LLM fine-tuning and concludes that SUPO supports training beyond context-length constraints. It identifies future work in advantage estimation, external memory, and cross-domain summarization optimization.
- SUPO integrates summarization into RL training and demonstrates strong empirical performance on CodeGym and BrowseComp-Plus against vanilla multi-turn RL.The framework jointly addresses tool-use behavior and summarization within standard RL infrastructure.
- Future work includes critic-based advantage estimation, external memory modules, and jointly optimizing summarization strategies across domains.
A.1 Proof of Theorem 3.2
The proof decomposes the policy gradient of a summarization-augmented rollout at states whose contexts exceed the summarization threshold. These indices partition the rollout into complete trajectories, yielding the stated gradient form.
- The policy gradient is obtained by differentiating the rollout objective and weighting log-policy derivatives by the terminal reward.
- Overlong states exceeding the summarization threshold identify the rollout indices used to decompose the policy-gradient summation.The proof uses states satisfying |s_h| ≥ L and containing the summary representation.
- The decomposition introduces h_0 = 0 and h_{I+1} = H, partitioning the MDP rollout into I + 1 complete trajectories.Each resulting segment shares the relevant prefix structure described in the proof.
- The proof concludes that the gradient can be expressed in the decomposed form after summing contributions across the identified trajectory segments.
B.1 Sample Problems
The appendix provides sample interactive tasks for CodeGym and BrowseComp-Plus, including their system prompts, tool interfaces, and example user problems. The examples emphasize stepwise tool use and explicit final-answer submission.
- CodeGym: CodeGym examples require agents to answer coding questions through provided function calls rather than writing code directly.The prompts allow at most one function call per step, require waiting for tool results, and use done to submit the answer.
- CodeGym: The CodeGym samples instantiate iterative tasks such as counting eligible student pairs under index and height constraints.The examples provide concrete arrays and function signatures for comparing student heights and submitting the count.
- CodeGym: CodeGym tools include observe for environmental information, problem-specific functions for intermediate operations, and done for final submission.The sample interfaces include functions such as compareHeights, calculateDelta, and findMaxDelta.
- BrowseComp-Plus: BrowseComp-Plus examples frame the agent as a research assistant that must conduct multi-step searches and produce evidence-supported answers.The system prompt emphasizes rigor, objectivity, thoroughness, and transparency.
- BrowseComp-Plus: BrowseComp-Plus tool calls use search and open_page for research, followed by finish for a concise answer and evidence-grounded explanation.The sample protocol specifies that document identifiers from search results can be used to open pages and cite evidence in the final response.
C More Experiment Results
The paper presents rollouts from CodeGym and BrowseComp-Plus to show the summarization patterns learned through end-to-end training. These patterns are examined together with the agents’ tool-use histories.
- Rollouts from CodeGym and BrowseComp-Plus illustrate summarization patterns learned through end-to-end training.The examples are intended to show how summarization interacts with tool-use history.
C.1.1 CodeGym
CodeGym rollouts improve from failing before training to succeeding after training, while summarization preserves task progress during long tool-use trajectories.
- Score 0 before training becomes score 1 after training on the CodeGym sample problem.
- Before training, the sample uses 4K/3K-token trajectories, whereas after training it uses 4K/2K-token trajectories.
- The learned summary retains the task objective, student heights, comparison rule, and current count of 28 so tool use can continue after context compression.
- After summarization, subsequent tool calls continue increasing the eligible-pair count from 28 to 43 before the agent submits the answer.
C.1.2 BrowseComp-Plus
BrowseComp-Plus trajectories improve from failure before training to success after training, with summaries preserving research objectives and unresolved verification steps.
- Both BrowseComp-Plus sample trajectories receive score 0 before training and score 1 after training.
- The post-summarization research state preserves the original identification question, verification checklist, and known findings about the international organization and Ghanaian university support.
C.2 Experiments for Scaling Trajectory Number
The scaling experiment evaluates accuracy under different test-time trajectory configurations, with full results reported in Table 3.
- Table 3 presents the full evaluation results for scaling the test-time number of trajectories.