FoldAct: Efficient and Stable Context Folding for Long-Horizon Search Agents

TL;DR

FoldAct tackles the non-stationary observation problem in long-horizon RL with context folding by separating gradient signals for summary and action tokens, enforcing full-context consistency via KL regularization between distributions under compressed and full contexts, and using selective segment training to reduce compute. The approach stabilizes training, preserves credit assignment between folding decisions and task actions, and yields substantial speedups (up to 5.19x) while maintaining competitive performance on Local RAG and Web Search benchmarks. Empirical results show improved accuracy and robustness against training collapse, with concise, information-rich summaries and effective context compression. These advances enable scalable, efficient long-horizon reasoning with LLM-based agents in complex, real-world information tasks.

Abstract

Long-horizon reinforcement learning (RL) for large language models faces critical scalability challenges from unbounded context growth, leading to context folding methods that compress interaction history during task execution. However, existing approaches treat summary actions as standard actions, overlooking that summaries fundamentally modify the agent's future observation space, creating a policy-dependent, non-stationary observation distribution that violates core RL assumptions. This introduces three fundamental challenges: (1) gradient dilution where summary tokens receive insufficient training signal, (2) self-conditioning where policy updates change summary distributions, creating a vicious cycle of training collapse, and (3) computational cost from processing unique contexts at each turn. We introduce \textbf{FoldAct}\footnote{https://github.com/SHAO-Jiaqi757/FoldAct}, a framework that explicitly addresses these challenges through three key innovations: separated loss computation for independent gradient signals on summary and action tokens, full context consistency loss to reduce distribution shift, and selective segment training to reduce computational cost. Our method enables stable training of long-horizon search agents with context folding, addressing the non-stationary observation problem while improving training efficiency with 5.19$\times$ speedup.

Figures (5)

• Figure 1: Comparison between standard RL (left) and RL with context folding (right). Left: Standard RL assumes stationary observations independent of the policy. Right Top: In context folding, summaries $\text{sum}_t \sim \pi_\theta(\cdot|s_t)$ are generated by the policy itself, and the visible state $s_t = [s_0, \text{sum}_t]$ is reconstructed from summaries, creating a self-conditioning loop where observations depend on $\pi_\theta$ (red warning). Right Bottom: This leads to training collapse: poor summaries from $\pi_{\theta_{\text{old}}}$ miss critical information, causing reward degradation and further policy degradation in a vicious cycle of non-stationary observations and unstable training.
• Figure 2: Our solution framework addresses all three challenges through three key components: (1) Separated loss computation (yellow/blue paths) computes independent gradients for summary and action tokens ($\mathcal{L}_{\text{summary}}$ and $\mathcal{L}_{\text{action}}$), preventing gradient dilution and enabling correct credit assignment (C1). (2) Consistency loss (green path) minimizes KL divergence between compressed context $s_t$ (used during rollout) and full context $h_{0:t}$ (stored offline), reducing distribution shift from self-conditioning (C2). (3) Selective segment training trains only on sampled turns per trajectory (C3).
• Figure 3: Summary Reward Design. The summary reward $r_{\text{summary}}(t)$ consists of two components: (1) Hallucination Penalty$r_{\text{hallucination}}(t)$ penalizes summaries generated when no information is available, preventing fabricated content. (2) Information Retention Reward$r_{\text{retention}}(t)$ rewards summaries that preserve critical information leading to task success. This design provides direct training signals for summary generation, enabling the agent to learn which information to retain in compressed contexts.
• Figure 4: Summary-Induced Context Variation and Selective segment Training. Left: During rollout, each turn $t$ sees a context $s_t$, which evolves as the agent progresses. As summaries are generated and inserted into the context (shown in blue), they replace earlier parts of the full history (shown in yellow). This results in different turns observing different contexts. Right: During training, PPO requires computing $\log \pi_\theta(a_t | s_t)$ for each sampled turn. Since each $s_t$ is unique—containing distinct summary and history content—each must be fed separately into the policy network $\pi_\theta$. This prevents reuse of intermediate computations (e.g., KV cache), making training costly. To reduce this cost, we apply selective segment training, where only a subset of turns (e.g., later ones) are used for loss computation. This reduces the number of context reconstructions and forward passes, lowering training cost without changing rollout behavior.
• Figure 5: Training stability analysis: (a) KL loss and (b) response length with and without consistency loss.