Technical Summary: TREK – Distill to Explore, Reinforce to Refine

1. Problem Statement

The paper addresses a specific failure mode in Group Relative Policy Optimization (GRPO) when applied to complex reasoning tasks. While GRPO is effective when the current policy (student) already samples useful solution trajectories, it often stalls on hard prompts where the correct solution modes lie outside the student's current on-policy support.

In these difficult regimes, the student repeatedly samples plausible but structurally similar incorrect trajectories. The scarcity is not the reward signal (verifiers can score correct solutions if found), but rather the exploration coverage of the solution space. Standard GRPO, being an on-policy method, cannot reinforce behaviors it does not sample. Consequently, larger rollout groups or sharper relative advantages often fail to escape the narrow region of the student's current support.

Existing distillation and teacher-informed methods largely focus on improving supervision or credit assignment for trajectories the student already reaches. They do not directly address the regime where the student rarely samples any reward-bearing trajectory. The core challenge is: How can we make useful behavior sampleable at all when the student's unaided exploration fails?

2. Methodology: TREK

The authors propose TREK (Teacher-Routed Exploration via Forward KL), a staged procedure that reframes distillation as exploration support expansion rather than simple imitation.

Core Philosophy

TREK treats a teacher (external or self-augmented) not as a universal behavior target, but as a proposal mechanism. The method relies on an output-only interface, meaning it requires only verified output trajectories, not teacher logits, hidden states, or probabilities. This allows the use of black-box teachers, white-box teachers, or the same model with additional inference-time context (e.g., reflection, longer reasoning budgets, or failure lessons).

The TREK Algorithm

The procedure separates prompt-level proposal availability from trajectory-level reachability:

1. Prompt Routing (Hard Prompt Mining):

   • For each prompt $x$, the system estimates the unaided student pass rate $p_S(x)$ by sampling $K$ rollouts.
   • If $p_S(x) \leq \tau_{low}$ (indicating a hard prompt where the student fails), the prompt is routed to the proposal phase. Easy prompts remain in the standard GRPO stream.

2. Proposal Generation & Selection:

   • A proposal source $\pi_T$ (e.g., a stronger model or the same model with extra context) generates candidate solutions for the routed hard prompts.
   • Trajectories are filtered by a verifier $V$. Only verified successes are kept.
   • Reachability Filtering: Among verified trajectories, the system ranks them by the current student's likelihood (measured via a trimmed, length-normalized Negative Log-Likelihood, $d_S$).
   • The top-$r$ student-proximal verified proposals are retained. This ensures the student is not asked to imitate solutions that are too far from its current capabilities, focusing instead on modes it can plausibly reach.

3. Forward-KL Consolidation (Warm-Start):

   • The retained trajectories form a dataset $D_{scaf}$.
   • A short Forward-KL update is applied: $\min KL(q_{prop} \parallel \pi_\theta)$.
   • This objective penalizes the student for assigning low probability to the verified proposal modes, effectively "pulling" these solution modes into the student's support. This is distinct from standard distillation (which often minimizes $KL(\pi_\theta \parallel q_{prop})$ or uses on-policy distillation).

4. On-Policy Refinement:

   • After the short Forward-KL warm-start, the prompts return to the standard GRPO rollout stream.
   • The student, now having expanded support to include the verified modes, can sample them and receive relative advantages from the verifier to refine the policy further.

3. Key Contributions

1. Identification of a Failure Mode: The paper identifies "missing exploration support" as a distinct bottleneck for GRPO on hard reasoning prompts, separate from reward sparsity.
2. Distillation as Exploration: It formulates a view of distillation specifically for expanding the student's reachable support, rather than just densifying supervision on already-sampled trajectories.
3. The TREK Procedure: A simple, routed algorithm that decouples prompt difficulty from trajectory reachability, utilizing a short Forward-KL phase to consolidate verified proposals before returning to GRPO.
4. Black-Box Compatibility: By requiring only verified outputs, the method is compatible with external black-box teachers, white-box teachers, or self-context augmentation, making it broadly applicable in practical settings.
5. Training Efficiency: It demonstrates that injecting verified solution modes before RL begins significantly accelerates learning on the hardest task types, reducing the number of optimization steps required to reach high performance.

4. Experimental Results

The authors evaluated TREK on mathematical reasoning (AIME 2024/2025) and agentic planning tasks (ALFWorld, ScienceWorld).

Mathematical Reasoning (Qwen3 Models)

• Setup: Tested on 1.7B, 8B, and 14B Qwen3 models. Proposals came from DeepSeek-V4 (external) or the same Qwen3 model with self-context (failure lessons).
• Performance: TREK consistently outperformed direct GRPO across all scales.
  • Qwen3-8B: Improved AIME 2025 from 36.9% to 40.3% and AIME 2024 from 47.9% to 51.1% (using DeepSeek-V4 proposals).
  • Self-Context Variant: Even without an external teacher, the self-context variant improved AIME 2025 to 38.5% and AIME 2024 to 49.6%.
• Ablation: Replacing the Forward-KL consolidation with On-Policy Distillation (OPD) style supervision resulted in lower performance, confirming that the Forward-KL objective is superior for expanding mode coverage on hard prompts.

Agentic Tasks (Qwen2.5-7B-Instruct)

• ALFWorld: Success rate increased from 75.8% (GRPO) to 82.8% (TREK with DeepSeek-V4) and 80.4% (self-context).
• ScienceWorld: Success rate more than doubled, rising from 12.5% (GRPO) to 26.7% (TREK with DeepSeek-V4) and 23.4% (self-context).
• Task Difficulty Analysis: Gains were concentrated on the hardest task types (e.g., "Heat & Place," "Examine in Light") where the baseline GRPO performance was lowest. Easy tasks showed little to no improvement, confirming the method targets the exploration bottleneck.
• Training Dynamics: TREK achieved high success rates much earlier in training compared to GRPO. For example, on ALFWorld, TREK exceeded 60% success by step 20, while GRPO remained below 50% and required ~5x more steps to reach the same level.

5. Significance and Claims

The paper claims that TREK offers a generalizable framework for improving on-policy reasoning systems when the primary bottleneck is the inability to sample reward-bearing trajectories.

• Complementarity: TREK is presented as complementary to existing methods (like token weighting or credit assignment improvements) which operate within the student's current support. TREK operates before that, ensuring the support exists.
• Efficiency: The primary advantage is not necessarily raising the asymptotic ceiling (as GRPO can eventually reach similar levels with massive compute), but drastically reducing the training budget required to discover hard solution modes.
• Practicality: The reliance on verified outputs rather than teacher internals makes the method robust for real-world deployment where black-box teachers or self-reflection loops are common, but internal model states are inaccessible.

The authors conclude that this "Distill to Explore, Reinforce to Refine" paradigm can extend to other post-training frameworks (e.g., DAPO, GSPO) and continual skill acquisition scenarios, providing a structured way to introduce new capabilities via verified proposals before stabilizing them through on-policy reinforcement.