Prefix-GRPO: метод обучения малых языковых моделей-агентов на префиксах траекторий

Researchers have proposed Prefix-GRPO — a reinforcement learning method that was published on arXiv (cs.LG) in July 2026. It helps small language model agents learn from the trajectories of a strong teacher model without turning them into one-off imitation, and across three environments — TextCraft, BabyAI, and ALFWorld — it outperforms both classical distillation and standard RL.

Why distillation stalls on long-horizon tasks

Direct distillation reduces the teacher's rich multi-step behavior to single-step targets for imitation — and this is inefficient wherever early decisions determine future states and rewards. Small models are appealing as a foundation for interactive agents: they are cheap and fast. But when an agent must go through a long chain of actions (pick up an item, reach a goal, follow an instruction), "repeat the teacher's answer" training loses the causal structure: the model memorizes the final utterances but doesn't understand how the teacher arrived at the required state.

How Prefix-GRPO works

Prefix-GRPO splits each teacher trajectory into two components: replay-aligned prefix queries and online continuations. Each prefix is replayed in the environment from scratch to restore a valid intermediate state, after which the student continues the interaction on its own and receives the reward for the task.

The key difference from response-only GRPO is that the method applies clipped policy updates to the historical assistant tokens inside the replayed prefix as well. The old log-probabilities of these tokens are estimated by a separate SFT checkpoint obtained through policy distillation. As a result, prefix training and continuation training are combined into a single form of policy optimization.

*Method: Prefix-GRPO, reinforcement learning built on top of GRPO

*Idea: the teacher's trajectory = a prefix for replay + an online continuation

*Difference: clipped updates apply to assistant tokens inside the prefix too

*Estimation of old log-probabilities — via a policy-distilled SFT checkpoint

*Test environments: TextCraft, BabyAI, ALFWorld (3 benchmarks)

What the experiments showed

Across all three environments, Prefix-GRPO improved agents on small models compared with both distillation and standard RL baselines. A separate series of ablation experiments checked exactly what drives the gain: replaying prefixes alone turned out to be insufficient without explicit optimization of the tokens inside them.

Replaying prefixes alone is not enough without explicit optimization

of the tokens inside the prefix.

— a statement from the study's abstract on arXiv. In other words, restoring the environment's intermediate state is not enough on its own — the policy also needs to be updated on the historical tokens, otherwise the method does not deliver its full effect. The implementation and the scripts for reproducing the results, according to the authors, have been released in open access on GitHub.

What this means

Prefix-GRPO is an attempt to transfer part of the "intelligence" of large teacher models into compact models without losing the multi-step logic of their behavior. If the approach scales beyond game-like environments, cheap local agents will be able to solve long tasks noticeably better than with ordinary distillation.

Frequently asked questions

What is Prefix-GRPO?

It is a reinforcement learning framework for small language model agents. It splits a teacher model's trajectories into prefixes, replays them in the environment to restore the intermediate state, and further trains the student in online interaction with a reward for the task.

How does Prefix-GRPO differ from ordinary GRPO?

Unlike response-only GRPO, which updates the policy only on new responses, Prefix-GRPO applies clipped updates to the historical assistant tokens inside the replayed prefix as well, estimating their old log-probabilities via a separate policy-distilled SFT checkpoint.

Where was the method tested?

The authors tested Prefix-GRPO on three agent environments — TextCraft, BabyAI, and ALFWorld — where it outperformed distillation and standard RL baselines. The code and reproduction scripts are available on GitHub.