ContextPilot: Teaching Agents for Proactive Context Management via Fine-grained RL

AI summary

Overview

Research area: Natural Language Processing — specifically long-horizon LLM agents, context management, and reinforcement learning (RL) training methods.

Technical level: Advanced. The paper assumes familiarity with agentic loop paradigms (e.g., ReAct), RL fine-tuning (GRPO), and credit assignment in multi-step trajectories.

Scope: ContextPilot is a proactive context-management framework that expands an agent's context-editing toolset and pairs it with a context-management-specific RL recipe using partial rollout and snapshot-level credit assignment.

Links given in the paper: project page at https://tencent.github.io/ContextPilot, source code at https://github.com/Tencent/ContextPilot, and models at https://huggingface.co/collections/panzs19/contextpilot. The paper is released under CC BY 4.0. Affiliations are Tsinghua University, Tencent Youtu Lab, and Shanghai AI Lab.

What This Paper Is About

Long-horizon agentic tasks require LLMs to iteratively retrieve, integrate, and maintain dispersed information across multi-turn interactions, but keeping every prior reasoning step, tool call, and tool response makes the working context grow without bound. Recent "proactive context management" methods let models edit their own working context with tools, but the paper identifies three limitations: a limited toolset (only search, deletion, and summarization, with no global planning, long-term memory, or adaptive compression), inefficient exploration that treats all context-management actions uniformly despite their heterogeneous impact, and coarse-grained credit assignment that spreads the final trajectory-level reward over all intermediate editing actions. ContextPilot is the authors' answer to all three.

Key Contributions

1. An extended context-management toolset. Building on the basic toolset of StateLM, ContextPilot adds tools for planning and perception (analyzeText, checkBudget, plan), long-term memory (memorize, updateMemory, readMemory), and soft context offloading (summarizeContext, compressContext, foldHistory). It also retains information retrieval tools (buildIndex, searchContext, readChunk, readMultiChunks) and note tools (note, updateNote, readNote), plus deleteContext from the original set.

2. Context-aware partial rollout. An RL training mechanism that uses context-length variation and entropy variation to score how "sensitive" each context-management action is, then spends extra branching/rollout budget on the top-ranked actions.

3. Fine-grained credit assignment. Instead of assigning the final trajectory-level reward to every intermediate context-editing action, the method estimates a snapshot's reward from all branched terminal trajectories that pass through that snapshot, and optimizes with GRPO treating each trajectory snapshot as an independent sample. Appendix A proves this estimator is unbiased and has variance σ²(S)/n_S versus σ²(S) for the trajectory-level estimator.

4. Extended evaluation. Whereas prior work mostly evaluates on long-context QA, this paper also evaluates on deep search tasks, showing consistent gains across multiple base models and benchmarks.

Main Findings

• Best average among comparable-size models with a much smaller window. ContextPilot uses only a 32K context window yet outperforms the 128K backbones. ContextPilot-8B-RL beats StateLM-8B-RL by an average of 3.55 points across four benchmarks; on deep search, ContextPilot surpasses SUPO by 1.51 average points across both backbones.

• Long-context QA numbers. ContextPilot-8B-RL scores 83.88 on NovelQA, 75.25 on ∞Bench, 64.27 on LongMemEval-S, and 54.18 on BrowseComp+ (average 69.40). ContextPilot-14B-RL reaches 84.81 / 81.08 / 67.40 / 55.50 (average 72.20). ContextPilot-E4B-RL reaches 72.92 / 60.99 / 62.47 / 47.47 (average 60.96).

• RL gains are larger on harder, longer-context tasks. ContextPilot-8B-RL improves over the SFT-only ContextPilot-8B by an average of 3.62 points across the four long-context QA tasks. NovelQA improvement is relatively modest because the SFT data already includes another split of NovelQA, whereas RL yields a 5.34 point increase on BrowseComp+. The paper notes NovelQA's average input length is about 119K tokens versus 552K tokens on BrowseComp+.

• Deep search numbers. With WebSailor-7B, ContextPilot scores 21.17 on BrowseComp (pass@3), 43.14 on BrowseComp-ZH (pass@3), 45.31 on GAIA (pass@1), and 43.67 on xBench-DS (pass@1), averaging 38.32. With WebExplorer-8B it scores 32.17 / 53.63 / 57.93 / 56.67, averaging 50.10.

• Token efficiency. Among trajectories with at least 15 turns, WebExplorer-8B's per-turn input length grows almost linearly on BrowseComp, reaching around 30K tokens, while ContextPilot-8B stabilizes per-turn input at roughly 8K–10K tokens.

• RL reshapes tool-use strategy. Early in RL training the model relies heavily on information retrieval tools, which account for roughly half of all tool calls; that share gradually decreases while planning/perception, long-term memory, and context offloading tools trend upward.

• RL improves tool-use correctness. Memory and context offloading tools show substantially higher invocation failure rates than information retrieval tools early in training, and proficiency with context management tools accompanies the improvement in task success rate.

• Tool-design ablation (Qwen3.5-397B-A17B). Cumulative tool additions raise average performance from 77.89 (original tools) to 80.29 (+planning) to 83.08 (+soft offloading) to 87.16 (+long-term memory). BrowseComp+ accuracy rises from 63.49% to 80.96% with the full toolset.

• RL-design ablation (Qwen3-8B). Entropy-based partial rollout improves several tasks but is unstable, decreasing BrowseComp+ accuracy by 1.32 points. Adding context variation yields more stable gains, and fine-grained credit assignment further improves over context-aware partial rollout across all four benchmarks (e.g., 75.25 on ∞Bench and 54.18 on BrowseComp+).

• Pilot study evidence. Branching from individual context-management actions and sampling 10 continuation rollouts each (evaluated with Qwen3-8B on NovelQA) shows that the standard deviation of final success rates differs substantially across tools, and a case study from StateLM-8B traces on BrowseComp+ shows final correctness can be misaligned with context-management quality.

Methodology in Plain English

The researchers start from the standard agent loop: at each step the model produces a thought and a tool call, the environment returns an observation, and all of it is appended to the context. They keep that loop but enrich the toolbox so the agent can do more than search, delete, and summarize — it can plan ahead, write structured long-term memory of entities, timestamps, and event episodes, and offload context softly (summarize, compress with a lightweight model such as llmlingua-2, or fold all history into keywords plus a summary that can be recovered later via keyword search).

To teach models to use this toolbox, they first build a "context management harness" that gives a teacher model (Qwen3.5-397B-A17B in thinking mode) step-appropriate hints and constraints — for instance, only exposing tools whose preconditions are met, and asking the teacher to retry on invalid arguments. Those scaffolding hints are stripped from the final supervised fine-tuning trajectories, which are then cut into snapshots at each context-editing operation, yielding 51,469 snapshots.

Then comes the RL stage, run with the verl library. Each query gets a budget of N=128 trajectory snapshots: trajectory-level rollouts first (8 rollouts, producing at most M=64 snapshots), and if the budget is not exhausted, the remaining slots go to partial rollouts branched at the highest-sensitivity context-management actions. Sensitivity combines normalized context-length change with entropy change measured against the trajectory's initial entropy, rather than against the preceding step, on the reasoning that partial rollout should find decisions that shift uncertainty relative to the initial query state rather than local fluctuations.

For credit assignment, the terminal snapshot receives an outcome reward, a format reward, and a penalty for invalid tool calls or context-length violations; each intermediate snapshot instead receives the average terminal reward over all sampled trajectories that pass through it. Snapshots from the same query are grouped, advantages are computed from the group mean and standard deviation, and GRPO optimizes the policy with each snapshot as an independent sample.

Why This Matters

Impact on research. The paper reframes context management as a first-class RL problem rather than a preprocessing heuristic, and shows that credit assignment at the level of individual context-editing decisions matters. Its proof that snapshot-level reward averaging is an unbiased, lower-variance estimator of conditional continuation value gives a concrete theoretical grounding for action-level credit in multi-step agent training. It also argues that evaluating on deep search, not only long-context QA, is necessary to see these effects.

Real-world applications:

• Enterprise and legal document analysis, where an agent must work through large corpora without letting the context balloon.
• Deep web research assistants that perform many searches and must decide what to keep, summarize, or discard.
• Long-running personal or task assistants that need persistent, structured memory across sessions.
• Cost-constrained deployments where a 32K-window model with good context management can substitute for a much larger-window model.

Industry relevance. Both a 32K context window replacing a 128K backbone and per-turn input stabilizing around 8K–10K tokens instead of growing toward 30K translate directly into lower inference cost and memory pressure for production agents. The paper also shows generalization across three long-context QA backbones (Qwen3-8B, Qwen3-14B, Gemma4-E4B-it) and two deep-search backbones (WebSailor-7B, WebExplorer-8B), which matters for teams that cannot retrain from scratch. Code and models are released.

Future Directions

• Richer context-editing operations. The authors state in the limitations that the toolset may still not cover all forms of context-editing demands, and suggest exploring more operations for organizing, compressing, and retrieving context under different task requirements.
• Hyperparameter search. Due to computational constraints the authors did not extensively search training hyperparameters, which they note may affect training efficiency and final performance, especially for partial rollout and credit assignment.
• Broader agentic settings. The paper's experiments focus mainly on long-context QA and deep search; extending proactive context management to agentic coding and GUI agents is named as an important future direction.
• Open question on sensitivity scoring. The ablation shows entropy-based partial rollout alone is unstable (a 1.32-point drop on BrowseComp+ for Qwen3-8B), leaving open how best to weight and combine the context-variation and entropy-variation signals.

Target Audience

Researchers and engineers working on LLM agents, long-context reasoning, and RL fine-tuning of language models — particularly those building multi-turn retrieval or search agents where context growth is the bottleneck, and those interested in fine-grained credit assignment for multi-step tool use. The paper is also useful to practitioners who want to substitute smaller context windows for larger ones without losing task performance.

Authors’ abstract

Long-horizon agentic tasks require large language models (LLMs) to iteratively retrieve, integrate, and maintain dispersed information across multi-turn interactions, but preserving all interaction histories leads to a continuously growing working context. Recent proactive context management methods allow models to edit their own working context with specialized tools, yet they still face three key limitations: (1) a limited toolset restricted to search, deletion, and summarization, with no support for global planning, long-term memory, and adaptive compression; (2) inefficient exploration that treats context management actions uniformly despite their heterogeneous impacts on final outcomes; and (3) coarse-grained credit assignment that assigns the final trajectory-level reward to all intermediate context editing actions during RL. To bridge these gaps, we introduce ContextPilot, a proactive context management framework for long-horizon agentic reasoning. Our approach systematically augments the toolset with planning, long-term memory, and soft context offloading tools. We further propose an RL method tailored for context management, which uses context and entropy variation to identify critical editing decisions for branch sampling and estimates action-level advantages from all branched trajectories that pass through the corresponding context editing action. Experiments on long-context QA and deep search tasks show that ContextPilot achieves stronger performance with a more compact working context, consistently outperforming existing baselines across various base models and benchmarks. Code is available at https://github.com/Tencent/ContextPilot.