From Trajectories to Prefixes: Reusing Teacher Trajectories via Replayed Prefixes and Online Continuation
Prefix-GRPO improves small models in long-horizon environments by reusing teacher trajectory prefixes and online continuation.
Key Findings
Methodology
Prefix-GRPO is a reinforcement learning framework that decomposes teacher trajectories into replay-aligned prefix queries and online continuations. Each prefix is replayed in the environment to recover a valid intermediate state, after which the student continues online interaction and receives task reward. Unlike response-only GRPO, Prefix-GRPO also applies clipped policy updates to historical assistant tokens inside the replayed prefix, using a policy-distilled SFT checkpoint to estimate their old log-probabilities.
Key Results
- In TextCraft, Prefix-GRPO achieved a Pass@8 of 96.00%, significantly outperforming Replay-GRPO's 78.79%.
- In BabyAI, Prefix-GRPO achieved an overall Pass@8 of 100.00%, demonstrating strong adaptability across diverse tasks.
- In ALFWorld, Prefix-GRPO achieved a Pass@8 of 94.50%, indicating its effectiveness in complex environments.
Significance
Prefix-GRPO addresses the inefficiency of small models in long-horizon interactive environments by decomposing teacher trajectories into prefix queries and online continuations. This method not only improves data efficiency but also enhances model adaptability in complex tasks, holding significant academic and industrial implications.
Technical Contribution
Prefix-GRPO uniquely unifies prefix learning and continuation learning within the same policy optimization form, significantly enhancing the learning efficiency of small models. By using a policy-distilled SFT checkpoint for clipped policy updates on historical assistant tokens, it opens new engineering possibilities.
Novelty
Prefix-GRPO is the first to decompose teacher trajectories into replayable prefix queries and online continuations, differing from traditional one-shot imitation targets by providing a more efficient learning path.
Limitations
- In some complex environments, Prefix-GRPO's average success rate remains lower than the strongest SFT-initialized RL baselines.
- Requires substantial computational resources for prefix optimization, possibly unsuitable for resource-constrained applications.
Future Work
Future research could explore more efficient prefix selection strategies and applications in resource-constrained environments. Further work could also investigate optimizing Prefix-GRPO's performance across different tasks and environments.
AI Executive Summary
Small language models are attractive for interactive agents due to their high-frequency deployment and low-latency interaction advantages. However, direct distillation from strong teacher trajectories often turns rich multi-turn behavior into one-shot imitation targets, which is inefficient in long-horizon environments. Prefix-GRPO addresses this by decomposing teacher trajectories into replay-aligned prefix queries and online continuations. Experimental results show that Prefix-GRPO significantly improves the performance of small-model agents in TextCraft, BabyAI, and ALFWorld, especially in complex tasks. Nevertheless, Prefix-GRPO's average success rate in some environments remains lower than the strongest SFT-initialized RL baselines. Future research could explore more efficient prefix selection strategies and applications in resource-constrained environments.
Deep Analysis
Background
In recent years, small language models have gained attention for their potential in interactive agents. Compared to large-scale models, small models offer significant advantages in high-frequency deployment and low-latency interaction. However, effectively transferring the rich decision structure from teacher trajectories to small models remains a challenge.
Core Problem
Direct distillation from strong teacher trajectories often simplifies complex multi-turn interactive behavior into one-shot imitation targets. This approach is inefficient in long-horizon environments because early decisions affect subsequent states and rewards.
Innovation
Prefix-GRPO provides a more efficient learning path by decomposing teacher trajectories into replay-aligned prefix queries and online continuations. This method not only improves data efficiency but also enhances model adaptability in complex tasks.
Methodology
- �� Decompose teacher trajectories into prefix queries and online continuations
- �� Use a policy-distilled SFT checkpoint to estimate old log-probabilities of historical assistant tokens
- �� Apply clipped policy updates to historical assistant tokens within replayed prefixes
Experiments
Experiments were conducted in three long-horizon interactive environments: TextCraft, BabyAI, and ALFWorld, evaluating Prefix-GRPO's performance across different tasks. Results show that Prefix-GRPO significantly improves the performance of small-model agents in all environments.
Results
In TextCraft, Prefix-GRPO achieved a Pass@8 of 96.00%, significantly outperforming Replay-GRPO's 78.79%. In BabyAI, Prefix-GRPO achieved an overall Pass@8 of 100.00%, demonstrating strong adaptability across diverse tasks. In ALFWorld, Prefix-GRPO achieved a Pass@8 of 94.50%, indicating its effectiveness in complex environments.
Applications
Prefix-GRPO can be applied in interactive agents requiring efficient decision-making, such as intelligent assistants, game AI, and automated customer service systems. Its excellent performance in long-horizon tasks makes it widely applicable in these fields.
Limitations & Outlook
Despite its strong performance across multiple environments, Prefix-GRPO's average success rate in some complex environments remains lower than the strongest SFT-initialized RL baselines. Additionally, prefix optimization requires substantial computational resources, possibly unsuitable for resource-constrained applications.
Plain Language Accessible to non-experts
Imagine you're cooking in a kitchen. You have an experienced chef as a teacher, guiding you through each step. Normally, you'd try to remember the entire process, which can be overwhelming. Prefix-GRPO is like breaking the chef's instructions into small segments, each a standalone step. This way, you can focus on the current step without worrying about the whole process. After each step, you receive feedback, allowing you to adjust your next move. This approach makes you more efficient and confident in cooking.
ELI14 Explained like you're 14
Imagine you're playing a complex video game. You have a super-skilled player guiding you on how to beat the game. Usually, you might try to remember all their tricks, but that's hard. Prefix-GRPO is like breaking their guidance into small segments, each a standalone task. This way, you can focus on the current task without worrying about the whole game. After each task, you get rewards, allowing you to adjust your strategy for the next move. This approach makes you more efficient and confident in gaming.
Glossary
Prefix-GRPO
A reinforcement learning framework that decomposes teacher trajectories into prefix queries and online continuations.
Used to enhance small models' performance in long-horizon environments.
SFT
A policy-distilled checkpoint used to estimate old log-probabilities of historical assistant tokens.
Used in Prefix-GRPO for prefix optimization.
Pass@8
The probability of succeeding at least once in eight attempts.
Used to evaluate model success rate in tasks.
Replay-GRPO
A reinforcement learning method that only replays prefixes.
Compared with Prefix-GRPO in experiments.
TextCraft
A long-horizon interactive environment for evaluating symbolic planning.
One of the environments used in Prefix-GRPO experiments.
Open Questions Unanswered questions from this research
- 1 How can Prefix-GRPO be efficiently applied in resource-constrained environments?
- 2 Can prefix selection strategies be further optimized to enhance learning efficiency?
Applications
Immediate Applications
Intelligent Assistants
Enhance decision-making efficiency in complex tasks using Prefix-GRPO.
Long-term Vision
Automated Customer Service Systems
Apply Prefix-GRPO to improve user interaction experience in automated customer service systems.
Abstract
Small language models are attractive backbones for interactive agents, but direct distillation from strong teacher trajectories often turns rich multi-turn behavior into one-shot imitation targets. This is inefficient in long-horizon environments, where early decisions shape later states and rewards. We propose Prefix-GRPO, a reinforcement learning framework that decomposes teacher trajectories into replay-aligned prefix queries and online continuations. Each prefix is replayed in the environment to recover a valid intermediate state, after which the student continues online interaction and receives task reward. Unlike response-only GRPO, Prefix-GRPO also applies clipped policy updates to historical assistant tokens inside the replayed prefix, using a policy-distilled SFT checkpoint to estimate their old log-probabilities. This unifies prefix learning and continuation learning within the same policy-optimization form. Experiments on TextCraft, BabyAI, and ALFWorld show that Prefix-GRPO improves small-model agents over distillation and standard RL baselines, while ablations show that replay alone is insufficient without explicit prefix-token optimization. The implementation and reproduction scripts are available at https://github.com/HappynessI/Prefix_GRPO.