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

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.

한 줄 요약

ContextPilot은 롱-호리즌 에이전트 작업에서 컨텍스트 관리를 강화하기 위해 RL 기반의 세부적인 프레임워크를 제안한 연구이다.

핵심 기여도

기존 툴셋(검색, 삭제, 요약)에 **글로벌 플래닝**, **장기 기억**, **소프트 컨텍스트 오프로딩** 툴을 추가.
RL 학습에서 **context-aware partial rollout**을 도입, **entropy variation**을 기반으로 중요한 컨텍스트 편집 결정을 샘플링.
**action-level advantage**를 추정하기 위해 모든 분기 경로에서 보상을 통합하여 **fine-grained credit assignment** 구현.
롱-컨텍스트 QA와 딥 서치 작업에서 기존 베이스라인 대비 **더 작은 컨텍스트 크기**로 **더 높은 성능** 달성.

핵심 아이디어

기존의 프로액티브 컨텍스트 관리 시스템은 툴셋이 제한적이고, RL 학습이 컨텍스트 편집의 영향을 세부적으로 반영하지 못해 효율성이 낮았다. ContextPilot은 이 문제를 해결하기 위해, **플래닝**, **장기 기억**, **소프트 오프로딩** 툴을 추가하여 컨텍스트를 보다 유연하게 관리하도록 설계했다. 또한, RL 학습 과정에서 **context-aware partial rollout**을 통해 컨텍스트 변화와 엔트로피 증가를 기준으로 중요한 편집 결정을 샘플링하고, **action-level advantage**를 추정하기 위해 모든 분기 경로의 보상을 통합함으로써 **fine-grained credit assignment**를 구현했다. 이는 기존의 trajectory-level reward만을 사용하는 방식보다 훨씬 정밀한 학습을 가능하게 한다.

기술적 접근법

**확장된 툴셋**: 기존의 검색, 삭제, 요약 외에 **글로벌 플래닝**, **장기 기억**, **소프트 컨텍스트 오프로딩** 툴 추가.
**RL 학습 방법**:
**context-aware partial rollout**: 컨텍스트와 엔트로피 변화량을 기준으로 중요한 편집 결정을 샘플링.
**fine-grained credit assignment**: 모든 분기 경로에서 보상을 통합하여 중간 편집 행동의 advantage를 추정.
**데이터셋**: 롱-컨텍스트 QA에는 NovelQA, ∞Bench, LongMemEval-S, BrowseComp+; 딥 서치에는 GAIA, BrowseComp, xBench-DeepSearch 사용.

주요 결과

**Long-context QA**에서 ContextPilot은 **더 작은 컨텍스트 크기**로 **더 높은 성능**을 보임.
예: NovelQA에서 기존 베이스라인 대비 **+7.2%** 개선.
∞Bench에서는 **+5.8%** 개선.
**Deep search** 작업에서도 **BrowseComp**에서 **+6.3%**, **xBench-DeepSearch**에서 **+4.9%** 개선.
모든 베이스 모델과 벤치마크에서 **일관된 성능 향상**을 기록.

의의 및 한계

ContextPilot은 롱-호리즌 에이전트 작업에서 컨텍스트 관리의 효율성과 정확도를 동시에 향상시킬 수 있는 새로운 프레임워크를 제시한다. 특히, **fine-grained credit assignment**와 **context-aware partial rollout**은 RL 기반 컨텍스트 관리의 새로운 기준을 제시한다. 그러나, **모든 툴이 동일한 영향을 미치지 않는다는 점**은 추가 연구가 필요하며, **다양한 도메인에서의 일반화 가능성**도 검증이 필요하다.

실용적 활용

ContextPilot은 **복잡한 다턴 대화 시스템**, **롱-컨텍스트 질문 답변**, **딥 서치 엔진** 등에서 유용하게 활용될 수 있다. 특히, **대규모 컨텍스트를 유지하면서도 효율적인 정보 관리가 필요한 산업**에서 실용적 가치가 높다.