Skip to content

TREK: Distill to Explore, Reinforce to Refine

🕒 Published (v1): 2026-07-06 00:00 UTC · Source: HuggingFace · link

Why this paper was selected

TREK fixes GRPO stalling on hard prompts via distillation-guided exploration — directly improves agent RL training

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

TREK (Teacher-Routed Exploration via Forward KL) addresses a core failure mode of GRPO—stalling on hard prompts where correct solutions lie outside the student's on-policy support—by using a proposal source (external or context-augmented self) to discover verified trajectories, then consolidating only the most student-proximal ones via a short forward-KL phase before returning to standard GRPO. On agentic benchmarks, TREK raises ALFWorld success from 75.8% to 82.8% and ScienceWorld from 12.5% to 26.7%, while also improving AIME 2024/2025 accuracy across all Qwen3 scales.

Problem

GRPO can only reinforce behaviors the student already samples. On hard prompts where the unaided student pass rate \(p_S(x) \approx 0\), the verifier is not the scarce resource—exploration coverage is. Larger rollout groups or sharper relative advantages still search within a narrow on-policy support, so the student repeatedly samples structurally similar wrong trajectories without discovering reward-bearing solution modes.

Method

TREK is a staged procedure:

  1. Hard prompt routing: Estimate unaided pass rate \(p_S(x) = \frac{1}{K}\sum_{i=1}^{K}\mathbf{1}[V(y_i^S)=1]\); route prompts with \(p_S(x) \leq \tau_{\text{low}}\) to proposal generation.
  2. Proposal selection: Query a proposal policy \(\pi_T\) (black-box teacher, white-box teacher, or same model with augmented context) for up to \(M\) rollouts; discard verifier-failing trajectories; rank survivors by a trimmed length-normalized student NLL (reachability score \(d_S(y_T|x)\)); keep the top-\(r\) student-proximal verified trajectories as \(\mathcal{Y}_{\text{reach}}(x)\).
  3. Forward-KL consolidation: On the accumulated scaffold dataset \(\mathcal{D}_{\text{scaf}}\), minimize \(\mathcal{L}_{\text{FKL}} = \text{KL}(q_{\text{prop}}(\cdot|x) \| \pi_\theta(\cdot|x))\), which reduces to teacher-forced NLL on retained proposals when teacher internals are unavailable: \(\mathcal{L}_{\text{prop}} = -\sum_t \log \pi_\theta(y_{T,t}|x, y_{T,<t})\).
  4. GRPO refinement: Return updated prompts to the ordinary GRPO stream.

The self-context variant augments the same deployment model with a failure-lesson memory synthesized from its own rejected attempts, avoiding any external teacher.

Key Contributions

  • Frames missing exploration support as a distinct failure mode of GRPO separable from reward sparsity or credit-assignment issues.
  • Proposes distillation as exploration support expansion (rather than imitation), using the proposal source purely as a verified trajectory generator—requiring only output trajectories, not teacher logits or hidden states, enabling black-box applicability.
  • Introduces the trimmed length-normalized student NLL \(d_S\) as a reachability score for selecting which verified proposals are close enough to the student's current support to be worth consolidating.
  • Demonstrates training acceleration on hard agentic tasks: TREK reaches high success rates in ~\(5\times\) fewer steps than unaided GRPO, which matters because agentic rollouts are expensive.

Results

Math reasoning (avg@16): - Qwen3-1.7B: AIME 2024 20.2 → 25.6 (+5.4); AIME 2025 17.5 → 20.3 (+2.8) with DeepSeek-V4 proposals - Qwen3-8B: AIME 2024 47.9 → 51.1 (+3.2); AIME 2025 36.9 → 40.3 (+3.4) - Qwen3-14B: AIME 2024 47.4 → 53.8 (+6.4); AIME 2025 39.7 → 44.6 (+4.9) - Self-context (no external teacher) Qwen3-8B: AIME 2024 49.6, AIME 2025 38.5 (vs. direct GRPO 47.9 / 36.9)

Agentic tasks (Qwen2.5-7B-Instruct, mean ± std over 10 runs): - ALFWorld: GRPO 75.8 ± 2.1 → TREK (DeepSeek-V4) 82.8 ± 2.7; self-context 80.4 ± 2.3 - ScienceWorld: GRPO 12.5 ± 1.7 → TREK (DeepSeek-V4) 26.7 ± 2.4; self-context 23.4 ± 2.1 - Hardest ALFWorld task types gain most: Heat & Place +19.1 pp, Examine in Light +12.6 pp - OPD (same proposals but on-policy distillation instead of forward-KL) trails TREK by 1.1–2.9 pp on AIME and underperforms on agentic tasks (ALFWorld 78.3, ScienceWorld 21.6) - TREK (self-context) exceeds 60% train success on ALFWorld by step 20; GRPO requires ~5× more steps to reach the same level

Limitations

  • Requires a proposal source capable of solving hard prompts that the student currently cannot; if the proposal source also fails on such prompts, no scaffold data is collected and TREK reduces to GRPO.
  • The self-context failure-lesson memory is heuristically designed (details in Appendix C, not fully specified in the paper); sensitivity to its quality is not thoroughly ablated.
  • MATH-500 gains are negligible because direct GRPO is near-saturated there—TREK's benefit is confined to genuinely hard exploration regimes.
  • Hyperparameters \(K\), \(M\), \(\tau_{\text{low}}\), \(r\), and the proposal-learning window length are reported only in an appendix; sensitivity to these is not deeply analyzed in the main text.
  • High run-to-run variance on agentic tasks (std up to ±2.7 pp) suggests evaluation is noisy; the margin over GRPO on ALFWorld (82.8 vs. 75.8) should be interpreted with this in mind.

Relevance to Agentic AI / LLM Agents

TREK directly targets long-horizon agentic tasks where hard-task exploration is the dominant bottleneck for RL fine-tuning—the ScienceWorld result (12.5% → 26.7%) shows the stakes are high. The paper's insight that a model can bootstrap itself via augmented-context self-proposals (verifier-guided retry, failure-lesson memory, backtracking) without an external teacher is practically significant because it removes the dependency on a larger or more capable model at training time. The training-efficiency angle—achieving strong performance on hard task types in a fraction of the steps unaided GRPO needs—matters specifically for agentic settings where each rollout involves costly environment interaction. This work connects the RL-from-verifiable-rewards line (GRPO, RLVR) with the self-improvement / iterative distillation literature in a way that is deployment-compatible.