PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhuoshi Pan, Qizhi Pei, Junru Lu, Honglin Lin, H. Vicky Zhao, Di Yin, Xing Sun

arXiv:2608.28476v1cs.CL

TL;DR

Long-horizon agents must manage rapidly growing interaction contexts, while existing tools and RL training provide limited planning, memory, adaptive compression, exploration, and credit assignment. ContextPilot expands context-management tools and introduces context-aware RL, achieving stronger performance with more compact contexts across long-context QA and deep search tasks.

  • Problem

    Existing context-management methods lack planning, long-term memory, adaptive compression, action-sensitive exploration, and fine-grained credit assignment for long-horizon tasks.

  • Method

    ContextPilot adds planning, long-term memory, and soft context offloading tools, then uses context-aware partial rollout and branched-trajectory advantages for RL training.

  • Results

    ContextPilot delivers consistent gains across long-context QA and deep search, with ContextPilot-8B-RL exceeding StateLM-8B-RL by 3.55 average points across four benchmarks.

  • Takeaways & Limitations

    ContextPilot shows that proactive context management can support stronger agentic reasoning while maintaining a more compact working context.

  • Takeaways & Limitations

    The toolset may not cover all context-editing demands, and experiments focus mainly on long-context QA and deep search tasks.

Abstract

from arXiv · show

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.

1 Introduction

Long-context QA and deep search require LLMs to identify evidence, integrate complex relations, and interact with tools across multiple turns. ContextPilot addresses limitations in existing context management by expanding the toolset and tailoring RL training to context-editing decisions.

  • Long-context QA and deep search [3] [4] require models to identify key evidence, integrate complex relations, and seek information through multi-turn tool interactions [5].
  • Existing context management tools centered on search, deletion, and summarization cannot support long-horizon needs for long-term memory [14] [15], structured entity-event episodes [16], and global planning [6] [17].
  • Context editing operations can substantially overwrite interaction history and have heterogeneous effects on later steps, with different operations producing substantially different final-success-rate variance.
  • ContextPilot also designs an RL training method tailored to context management rather than relying solely on traditional trajectory-level rollouts.
  • ContextPilot extends context management with planning, long-term memory, and soft context offloading tools, enabling better control over context during long-horizon interactions.

2 Related Work

Prior work managed growing interaction histories through predefined truncation, summarization, compression, or extraction workflows, leaving models passive recipients of externally managed context. More recent approaches equip agents with tools to actively edit working context and memory, including virtual-memory management and learned memory editing.

  • Passive Context Management: Rule-based methods [8] [10] [21] reduce context length through fixed workflows or external modules, but leave models passive recipients of context management.These methods truncate or summarize messages after length thresholds, or compress and extract information from historical context.
  • Active Context Management: Recent agentic methods let models decide when and how to edit working context through context-management tools.MemGPT [23] treats agent memory as virtual memory managed through tool calls, while Yu et al. [18] and Xu et al. [15] train agents to control memory without an external controller.

3 Preliminary

This section formalizes iterative agentic reasoning and proactive context management, whose appended interactions cause context growth. It also reviews snapshot-based training and pilot evidence that context-management actions differ substantially in their impact on final outcomes.

  • Context management preliminaries: Agentic reasoning iteratively generates thoughts and tool calls from the current context, appending environment feedback so the context grows monotonically.At step i, the model produces thought t_i and tool call a_i from context c_i; feedback is appended to form the next context.
  • Context management preliminaries: Proactive context management lets agents manipulate their accumulated context through specialized context-management actions.These actions update the current context through a context transition function.
  • Training context management: Existing training addresses intermediate context edits by segmenting trajectories into independent snapshots and masking previously seen outputs.For M context-editing operations, a trajectory is divided into M + 1 snapshots, ensuring intermediate states contribute to training without redundant optimization.
  • Pilot study: Branch-sampled continuations show that context-management actions have substantially different effects on final success rates, motivating greater exploration for high-impact actions.The pilot study branches from context-management actions in existing trajectories and measures the standard deviation of continuation success rates.

4 Methodology

ContextPilot extends context management with planning, structured long-term memory, and soft offloading tools, then trains agents using adaptive partial rollouts and snapshot-level credit assignment. Its reinforcement-learning recipe concentrates exploration and reward estimation on consequential context-editing decisions.

  • Extended Toolset: ContextPilot adds planning, structured long-term memory, and soft context offloading to StateLM’s basic toolset.The memorize tool extracts entities, timestamps, and event episodes while building links among memory items; readMemory retrieves a target memory with its neighbors.
  • SFT Data Construction: A context-management harness guides SFT data generation with step-specific hints, including inspecting search results and offloading context after length thresholds are exceeded.These hints and constraints are used only during teacher-data generation.
  • Context-Aware Partial Rollout: Context-aware partial rollout ranks context-editing actions by combined context and entropy variation, allocating extra branches to the most sensitive decisions.Entropy is referenced to the initial query state to capture substantial uncertainty changes rather than local fluctuations; unused rollout budget is assigned to top-ranked actions.
  • Fine-Grained Credit Assignment: Fine-grained credit assignment estimates each intermediate snapshot from all terminal trajectories sharing that snapshot as a prefix, rather than copying one final reward across actions.Rewards combine answer correctness, output parsability, and penalties for invalid tool calls or context-length violations.
  • Fine-Grained Credit Assignment: GRPO [28] optimizes the policy with trajectory snapshots as independent samples, shifting credit assignment from complete trajectories to intermediate context-management decisions.The paper states that this enables more precise reward estimation and discusses variance reduction in Appendix A.

5 Experiments

Experiments on long-context QA and deep search show that ContextPilot improves benchmark performance while maintaining a compact working context. Ablations further support its expanded toolset and context-aware, fine-grained RL design.

  • Main results: ContextPilot achieves the best average performance among comparable-size models, using a 32K context window while outperforming a 128K backbone and prior RL-trained agents.ContextPilot-8B-RL outperforms StateLM-8B-RL by an average of 3.55 points across four long-context QA benchmarks and surpasses SUPO by 1.51 average points across both backbones.
  • RL training: RL further improves ContextPilot over SFT, with ContextPilot-8B-RL gaining 3.62 average points across four long-context QA tasks.The gain on NovelQA is relatively modest because SFT includes another NovelQA split, whereas RL produces a substantial 5.34-point improvement on another reported setting.
  • Token efficiency: ContextPilot stabilizes BrowseComp input length at roughly 8K–10K tokens per turn, whereas WebExplorer-8B grows almost linearly to around 30K tokens.The token-efficiency analysis considers trajectories with at least 15 turns and computes average input tokens per turn.
  • RL behavior: During RL, retrieval tool use decreases while planning, perception, and long-term memory use increase, alongside improved tool-use correctness and task success.Early training relies on retrieval for roughly half of tool calls; tool invocation failures include malformed calls, invalid arguments, and violated tool preconditions.
  • Tool design: Adding planning, soft context offloading, and long-term memory tools progressively improves performance, with BrowseComp+ accuracy rising from 63.49% to 80.96%.The full cumulative toolset achieves the highest score in the tool-design ablation.
  • RL training design: Context variation stabilizes partial rollout gains, whereas entropy alone can reduce BrowseComp+ accuracy by 1.32 points; fine-grained credit assignment improves further.The RL ablation evaluates context-aware partial rollout and fine-grained credit assignment.

6 Conclusion

ContextPilot combines an expanded proactive context-management toolset with context-aware partial rollout and fine-grained credit assignment for RL training. Experiments show stronger performance with a more compact working context on long-context QA and deep search tasks.

  • Conclusion: ContextPilot achieves stronger performance while maintaining a more compact working context on long-context QA and deep search tasks.The gains are reported consistently across both task types.
  • Conclusion: The system extends context management with planning, long-term memory, and soft context offloading tools.These additions broaden the toolset beyond basic context editing operations.
  • Conclusion: Its RL training uses context-aware partial rollout and fine-grained credit assignment to improve context-management decisions.The ablation terminology distinguishes entropy-based partial rollout, context-aware partial rollout, and fine-grained credit assignment.

Limitations

ContextPilot’s toolset may not cover all context-editing demands, and computational constraints limited extensive exploration of some training hyperparameters.

  • Toolset limitations: The current toolset may not support every context-editing demand, motivating richer operations for organizing, compressing, and retrieving context across task requirements.Future work should broaden context-management operations to handle diverse task-specific needs.
  • Training limitations: Computational constraints prevented extensive searches over some training hyperparameters.

A Theoretical Discussion of Fine-Grained Credit Assignment · B Training Dataset Statistics

ContextPilot’s fine-grained credit estimator averages terminal rewards across sampled continuations from the same context snapshot, yielding lower variance and more stable optimization signals under stated sampling assumptions. Training uses filtered SFT trajectories from NovelQA and NarrativeQA, segmented into 51,469 context-editing snapshots, alongside LongBench-v2 data for RL.

  • A Theoretical Discussion of Fine-Grained Credit Assignment: The theoretical comparison conditions on continuations sampled independently from the same distribution p(· | S), with nS denoting the number of sampled continuations.Under these conditions, averaging all continuations passing through S preserves unbiasedness while reducing variance relative to using one continuation.
  • A Theoretical Discussion of Fine-Grained Credit Assignment: The proposed estimator targets the same conditional continuation value Q(S) as trajectory-level credit assignment while reducing estimator variance.The result assumes independent continuation rewards with finite conditional variance and follows because both estimators are unbiased.
  • A Theoretical Discussion of Fine-Grained Credit Assignment: Lower-variance reward estimates provide GRPO with a more stable relative-advantage signal for optimizing context-editing decisions.Context-aware partial rollout strengthens this effect by increasing nS for sensitive context-management actions.
  • B Training Dataset Statistics: SFT training retains quality-filtered trajectories from NovelQA’s PublicDomain split and NarrativeQA’s training split, then segments them at context-editing operations.These operations produce 51,469 trajectory snapshots in total.
  • B Training Dataset Statistics: RL training uses 488 LongBench-v2 quThe supplied passage truncates the remainder of the RL dataset description.

C SFT Data Synthesis Details

SFT trajectories are synthesized for long-context QA using a guided Qwen3.5-397B-A17B teacher and filtered for answer correctness, context-management quality, and context length. The process combines dynamic tool presentation with retry-based outcome filtering and process-based screening.

  • Trajectory generation: SFT trajectories target long-context QA and are generated by Qwen3.5-397B-A17B in thinking mode with temperature 0.6, top-p 0.95, and a 4K-token output limit.A context-management harness supplies tool definitions and procedural guidance for handling long interaction histories.
  • Trajectory generation: The harness dynamically presents context-management tools rather than exposing the full toolset at every step.
  • Trajectory filtering: A three-stage pipeline retains initially incorrect samples when either of two additional retries produces the correct answer, removes improper context-management behavior using GPT-OSS-120B [44], and discards trajectories exceeding 32K peak context tokens.Outcome filtering uses exact match before process-based filtering and the peak-context-length constraint.

D Training Details … G Tool Failure Definitions

The paper specifies training protocols, benchmark datasets, baseline-reproduction settings, and evaluation procedures for context efficiency and tool failures. ContextPilot uses proactive context management to keep working contexts compact while defining tool failures by environment-side execution errors.

  • D Training Details: Training data covers NovelQA, NarrativeQA, and LongBench-v2 for long-context QA and OpenSeeker-v1 for deep search.These dataset categories are summarized in Table 6.
  • D Training Details: ContextPilot trains with SFT and RL using verl [36], with ZeRO-3 SFT and 128-step RL using rollout batches of 16 and KL coefficient 0.001.SFT uses global batch size 128, learning rate 5 × 10−6, cosine scheduling, and 0.03 warmup; RL samples 8 rollouts per query and segments trajectories into up to 8 snapshots.
  • D Training Details: ContextPilot uses task-specific system prompts for long-context QA and deep search, while GPT-OSS-120B [44] grades open-ended outputs with the StateLM prompt [11].The QA and deep-search prompts appear in Figures 7 and 8, and the judging prompt appears in Figure 9.
  • E Baseline Reproduction Details: ReSum [8] is reproduced without training using WebSailor-7B and WebExplorer-8B, with a 32K context window, 25K-token summarization trigger, and 60-call tool budget.The setting allocates 30K input tokens and 2K generation tokens; Qwen3-30B-A3B is used because the official summarization model is unavailable.
  • E Baseline Reproduction Details: SUPO [9] is trained on the same 1K OpenSeeker samples as ContextPilot, jointly optimizing policy-generated summaries and tool-use actions.Its summarization threshold is L = 0.95 × 30K = 28.5K tokens, with maximum input length 30K.
  • F Additional Token Efficiency Analysis: Across BrowseComp and BrowseComp-ZH, ContextPilot keeps working contexts substantially more compact than WebExplorer-8B as interaction turns increase.The analysis considers trajectories with at least 15 interaction turns and reports average input tokens per turn.
  • G Tool Failure Definitions: A tool invocation fails when execution triggers an environment-side error, including malformed calls, invalid arguments, or violated tool-specific preconditions.For retrieval, no matched evidence is not a failure when the query and index are valid; Table 7 and Figure 6 summarize cases and category rates.

H Additional RL Training Ablation · I Use of Scientific Artifacts · J Use of AI Assistants

The RL ablations show that context-aware partial rollout and fine-grained credit assignment progressively improve average performance across Qwen3-8B and Gemma4-E4B-it. The paper also documents scientific-artifact licensing and responsible AI-assistant use, including manual verification of assisted outputs.

  • H Additional RL Training Ablation: Context-aware partial rollout improves over entropy-only branching, while fine-grained credit assignment further strengthens average performance on both evaluated models.The complete ablation covers Qwen3-8B and Gemma4-E4B-it.
  • I Use of Scientific Artifacts: The paper cites the original creators of scientific artifacts wherever those artifacts are introduced.
  • I Use of Scientific Artifacts: Qwen3 and Gemma4 use Apache License 2.0, OpenSeeker-V1 uses MIT, and the other datasets use Apache License 2.0.All artifacts are used for research consistently with their intended benchmarking, training, and evaluation purposes, without redistributing the original datasets.
  • J Use of AI Assistants: Codex assisted code modification, while ChatGPT assisted writing polishing.The authors manually reviewed, verified, and finalized all assisted outputs.
  • J Use of AI Assistants: The paper includes a tool execution failure reason analysis in Table 7.
  • H Additional RL Training Ablation: Table 8 reports the complete RL training ablation results for Qwen3-8B and Gemma4-E4B-it, with parenthetical values denoting changes from the preceding line.
Loading 2608.28476v1…