1 Introduction↩︎

Endowing small language models with robust agentic capabilities has emerged as a central challenge, as these models transition from static question answering to interactive decision-making in real-world environments. Compared with frontier-scale models, small models offer compelling advantages for high-frequency deployment, low-latency interaction, and private or on-device inference, making them promising backbones for specialized agents [1]. However, multi-turn agent tasks are substantially harder than static generation: each action changes the future observation distribution, so the model must maintain coherent state tracking, planning, action execution, and stopping behavior throughout a closed observation-action-feedback loop [2], [3].

Using strong teacher agents to generate high-quality interaction trajectories and then transfer them to small models is a natural approach [4]–[6]. Yet direct supervised distillation over full trajectories often teaches the surface form of teacher behavior rather than the underlying decision dynamics needed for long-horizon interaction. As a result, distilled small models may exhibit verbose reasoning, excessive tool use, unstable stopping, or other forms of inefficient behavior without learning which historical decisions are actually responsible for downstream success [7]–[9]. This creates a basic mismatch: teacher trajectories contain rich sequential decision structure, but standard distillation typically consumes them as one-shot imitation targets.

Recent work has begun to explore how demonstrations, historical prefixes, and off-policy trajectories can improve the data efficiency of reinforcement learning. Some methods use prefixes as rollout conditions or starting states [10], [11], while others emphasize the influence of early reasoning on downstream trajectories [12]–[14]. However, these methods still leave two issues unresolved for reusing teacher prefixes in RL: how to split one long teacher trajectory into multiple valid training states, and how to avoid consuming distilled trajectories as one-shot demonstrations or passive conditions, which underuses the decision information distributed across the trajectory.

Prefix-GRPO addresses this inefficiency by treating trajectory splitting as a way to create differentiated learning queries from a single teacher rollout. Rather than using the whole trajectory as one supervised target, our group-query construction selects informative cut states along the teacher trajectory, so that different prefixes expose different partial progress, remaining subgoals, and recovery opportunities. Each query pairs a teacher prefix that reconstructs an intermediate state with a student continuation optimized by online reward. We further bring the assistant tokens in the teacher prefix into the same clipped policy-optimization form as the continuation tokens, using an SFT-based old-policy estimate for the prefix side. In this way, Prefix-GRPO improves the utilization of distilled trajectories: one teacher trajectory supplies multiple decision-aware training queries, and both sides of each cut contribute to policy learning.

Our main contributions are as follows:

• We propose Prefix-GRPO, a reinforcement learning framework for multi-turn agents that jointly learns from replayed teacher prefixes and online continuations, enabling teacher trajectories to be reused beyond one-shot distillation.

• We introduce a replayable group-query construction that turns each teacher trajectory into multiple cut-state training queries at different historical depths.

• We develop a prompt-space prefix old-logprob anchoring mechanism that enables historical prefix tokens to participate in a clipped GRPO-style objective.

2 Related Work↩︎

Recent research has re-emphasized the cost, latency, and deployment advantages of small models in agent systems [1]. Recent work on small-model reasoning further shows that controlled exploration and refined offline integration can improve SLM training when distillation is combined with online RL [15]. In parallel, several studies investigate how to transfer tool use, retrieval behavior, and multi-step interaction abilities from large-model agents to smaller students [4]–[6]. These methods usually treat teacher behavior as imitation targets, in-context demonstrations, or reusable agent modules. Our setting also focuses on small-model agents, but the core question is different: we study how a teacher interaction trajectory can be decomposed into prefix queries that can be replayed in the environment, rather than how to compress the teacher’s behavior in one shot.

More broadly, recent studies on graph-aware and structural-data adaptation show that language models often require explicit mechanisms to absorb relational structure, including attention analyses over graph-structured inputs, graph-aware recommendation learning, and structural graph wordification [16]–[18]. This perspective is complementary to our setting, where the structure to be reused comes from temporal interaction histories rather than static graphs.

Another line of work studies result-driven reinforcement learning for language models, especially GRPO- or RLVR-style optimization in mathematics, code, and other verifiable domains [19]–[21]. These methods typically optimize response tokens generated after the current prompt, leaving prompt-history assistant tokens outside the clipped objective because their old-policy log-probabilities are not directly available from online rollout. Recent analysis of asynchronous agentic RL further shows that missing old logits can create semantic mismatch in off-policy correction when the reference policy is not explicit [22]. In parallel, multi-turn agent RL places language models inside observation-action-feedback loops and studies environment interfaces, long-horizon rollout, tool use, and training stability [2], [3], [23]–[27]. Prefix-GRPO inherits the PPO/GRPO-style ratio clipping, but extends the optimization target from continuation tokens to assistant prefix tokens from teacher interaction history.

The most closely related studies use demonstrations, historical prefixes, or off-policy trajectories to improve RL sample efficiency [10], [11], [14]. These methods typically use prefixes as rollout conditions, starting states, or difficulty-control signals, guiding the policy toward local regions that are more likely to yield learning signal. Other work emphasizes that prefixes themselves can strongly influence downstream reasoning paths or optimization stability [12], [28]. Prefix-GRPO shares the view that prefixes are central to learning, but focuses on a different use of trajectory structure: it explicitly splits teacher interaction trajectories into multiple prefix queries that can be replayed in the environment, and treats the assistant tokens inside those prefixes as optimization targets rather than only as conditioning context.

3 Method↩︎

We start from a collection of multi-turn teacher trajectories. In our experiments, these trajectories are generated by MiniMax-M2.1 across the evaluated environments, with each environment instance sampled multiple times to obtain diverse successful interaction traces; the algorithm itself does not depend on this particular teacher. Let \(\tau_i=(h_{i,1},\ldots,h_{i,T_i})\) denote the \(i\)-th teacher trajectory, where each \(h_{i,t}\) is a serialized message or environment-feedback element. Unlike conventional response-only RL, which usually treats each trajectory as a single imitation target or extracts only one cut point, Prefix-GRPO expands one teacher trajectory into multiple cut-state queries and optimizes both the continuation after the cut state and the historical assistant prefix that leads to it.

3.1 Prefix Construction with Group Queries↩︎

In this subsection, we describe how Prefix-GRPO derives multiple prefix queries that can be replayed in the environment from each teacher trajectory and turns them into aligned cut-state training samples.

3.1.0.1 Query selection.

For each trajectory \(\tau_i\), group-query construction selects a set of cut-state queries \(Q_i=\{q_{i,1},\ldots,q_{i,J_i}\}\), where \(J_i\) is the number of retained queries and each query corresponds to a teacher prefix ending immediately after an assistant turn. Figure 2 illustrates the construction. We first score candidate token locations on the serialized teacher trajectory using teacher-forcing entropy under the policy-distilled SFT model. The score is restricted to assistant interaction tokens, smoothed with a centered window, and converted into an entropy-change signal. High-change tokens are then mapped back to their assistant turns, so each selected token induces a cut point after the corresponding assistant action. In the example, the top-ranked entropy-change tokens produce three cut-state queries at different historical depths; each query becomes a prompt of the form system instruction, prefix history, and cut-state observation. We keep up to three distinct prefix queries per trajectory, together with one raw no-prefix query. Exact scoring, deduplication, and top-\(k\) rules are given in Appendix 6.2.0.1.

3.1.0.2 Replay validation and canonicalization.

Each selected query is replayed in the environment before it is admitted into training. In TextCraft, replay starts from an environment created with the sample’s task identifier and goal, executes the teacher prefix actions sequentially, and keeps the query only if the replayed observations match the continuation-side observations on shared structured fields. Queries marked mismatch, unverifiable, or error are discarded rather than repaired.

After replay, the training prompt is canonicalized as \[\label{eq:canonical-prompt} x_i=\mathrm{concat}(\text{system}, \text{prefix history}, \text{cut-state user observation}).\tag{1}\] This canonical form makes the prompt end at the replayed cut-state observation rather than at the last teacher action, ensuring that student rollout starts from an environment-consistent state. Each replay-validated query is then treated as an independent training sample; for notational simplicity, we re-index these samples by \(i\) below.

3.2 Old-Logit Estimation for Distilled Prefixes↩︎

In this subsection, we describe how Prefix-GRPO computes old log-probabilities for distilled prefix tokens in the canonicalized prompt space.

3.2.0.1 Prompt-space old-policy estimation.

Continuation tokens are sampled by the online rollout policy, so their old log-probabilities are available in the standard GRPO buffer. Prefix tokens are different: they are teacher-derived assistant tokens already embedded in the prompt history. To place them under the same clipped-ratio training form, we use the policy-distilled SFT checkpoint as a fixed old policy for prefixes and recompute teacher-forcing log-probabilities on the canonicalized prompt. Concretely, each replay-validated sample stores three prompt-space annotations: (i) an assistant-prefix span \([b_i,e_i)\), (ii) a binary prefix mask \(m_i\) over that span, and (iii) cached old log-probabilities \(\ell_{i,u}^{\mathrm{old}}\) stored densely over the same span. The span localizes the replayed prefix window, while the mask specifies which positions inside that window should receive prefix loss. This is necessary because the window is defined in prompt-token coordinates and may contain boundary, formatting, or non-optimized tokens; the objective should update only the historical assistant tokens selected for prefix learning. The valid prefix-token set is therefore \[\label{eq:prefix-mask-set} M_i^{\mathrm{pre}}=\{u\mid b_i\le u<e_i,\; m_i[u-b_i]=1\}.\tag{2}\] This representation is deliberately prompt-space rather than turn-space: after canonicalization, the prefix tokens may lie in the middle of the prompt, and the span-mask-logprob annotations provide the coordinates needed to gather them during training. More detail is given in Appendix 6.2.0.4.

3.3 Unified Prefix-Continuation Optimization↩︎

In this subsection, we describe how Prefix-GRPO writes continuation learning and prefix learning in a unified clipped-ratio form.

3.3.0.1 Token groups.

Prefix-GRPO optimizes two token groups in the same policy update. The continuation group contains tokens generated by the student after the replayed cut state and is trained with environment reward. The prefix group contains historical assistant tokens from the replayed teacher prefix and is trained with the cached prompt-space old log-probabilities described above. The two groups differ in where their tokens come from and how their old log-probabilities are obtained, but both are optimized by the same clipped surrogate.

3.3.0.2 Continuation branch.

Let \(M_i^{\mathrm{cont}}\) be the set of valid continuation tokens and let \(\hat{A}_{i,t}^{\mathrm{cont}}\) denote the token-level continuation advantage, obtained by broadcasting the sequence/group-level advantage to valid continuation positions. The continuation loss is \[\label{eq:cont-loss} \mathcal{L}_{\mathrm{cont}} =-\sum_{t\in M_i^{\mathrm{cont}}} \min\left( r_{i,t}^{\mathrm{cont}}\hat{A}_{i,t}^{\mathrm{cont}}, \mathrm{clip}(r_{i,t}^{\mathrm{cont}},1-\epsilon,1+\epsilon)\hat{A}_{i,t}^{\mathrm{cont}} \right).\tag{3}\] This branch is close to standard response-only GRPO. The main difference is not the surrogate form itself, but the fact that rollout starts from a replayed intermediate state rather than from the environment’s initial state.

3.3.0.3 Unified clipped surrogate.

For any valid prefix position \(u\in M_i^{\mathrm{pre}}\), the prefix importance ratio is \[\label{eq:prefix-ratio} r_{i,u}^{\mathrm{pre}} =\exp\left( \log \pi_\theta(z_{i,u}\mid x_i,z_{i,<u})-\ell_{i,u}^{\mathrm{old}} \right).\tag{4}\] For continuation positions, \(r_{i,t}^{\mathrm{cont}}\) is defined analogously with the rollout old policy. Both token groups are then optimized by clipped policy-ratio terms. The continuation term uses the group-relative advantage \(\hat{A}_{i,t}^{\mathrm{cont}}\) from online rollout. For the prefix group, Prefix-GRPO supports several choices of \(\hat{A}_i^{\mathrm{pre}}\); our main TextCraft configuration uses the continuation-magnitude anchor: \[\label{eq:prefix-adv} \hat{A}_i^{\mathrm{pre}} =\left| \frac{\sum_t \hat{A}_{i,t}^{\mathrm{cont}} m_{i,t}^{\mathrm{cont}}}{\sum_t m_{i,t}^{\mathrm{cont}}} \right|,\tag{5}\] where \(m_{i,t}^{\mathrm{cont}}\) is the valid continuation-token mask. For compactness, define \[g(r,A;\epsilon^-,\epsilon^+) =\min\!\left(rA,\mathrm{clip}(r,1-\epsilon^-,1+\epsilon^+)A\right).\] The complete training objective is \[\label{eq:full-loss} \begin{align} \mathcal{L} =&-\sum_{t\in M_i^{\mathrm{cont}}} g(r_{i,t}^{\mathrm{cont}},\hat{A}_{i,t}^{\mathrm{cont}};\epsilon,\epsilon) \\ &-\lambda_{\mathrm{pre}}\sum_{u\in M_i^{\mathrm{pre}}} g(r_{i,u}^{\mathrm{pre}},\hat{A}_i^{\mathrm{pre}}; \epsilon_{\mathrm{pre}}^{-},\epsilon_{\mathrm{pre}}^{+}). \end{align}\tag{6}\] This form puts prefix learning and continuation learning under the same clipped surrogate, while keeping their loss terms, old-policy sources, and advantage assignments separate. It should not be interpreted as assigning signed credit to the replayed teacher prefix. Signed credit remains in the continuation term, while the prefix term uses the magnitude of the continuation learning signal to control how strongly replay-validated teacher prefixes are anchored. The coefficient \(\lambda_{\mathrm{pre}}\) controls the strength of prefix optimization. The main experiments use the split objective shown above, while joint-objective variants are studied as ablations. Additional objective sensitivities are reported in Appendix 8.1.

4 Experiments↩︎

4.1 Experimental Setup↩︎

4.1.0.1 Environments and tasks.

We evaluate on three long-horizon interactive environments: TextCraft for symbolic crafting [29], ALFWorld for household instruction following [30], and BabyAI for grid-world compositional instruction execution [31]. Together, these environments cover symbolic planning, grounded instruction following, and long-horizon action execution, allowing us to evaluate whether Prefix-GRPO improves small-model agents across different forms of multi-turn decision making.

4.1.0.2 Baselines.

We compare five method families: the raw Base Model, SFT as direct trajectory distillation, RL recipes from the base model (GRPO, GRPO-MIS, DAPO) [19], [20], the same RL recipes from an SFT-initialized policy, and Prefix-Based RL, which includes Replay-GRPO and Prefix-GRPO. The key contrast is whether teacher prefixes are ignored, replayed only, or replayed and explicitly optimized.

4.1.0.3 Environment-specific settings.

Each environment uses its official task-family split and action interface. We keep the evaluation budget fixed at eight sampled rollouts per task and report both Avg@8 and Pass@8. Detailed prompt formats, SFT configurations, rollout limits, and Prefix-GRPO training settings for each environment are provided in Appendix 7.

4.2 Effectiveness of Prefix-GRPO↩︎

All experiments use Qwen3-1.7B as the base policy. For each task, we sample eight rollouts and report two task-family metrics: Avg@8, the average success rate across the eight rollouts, and Pass@8, whether at least one rollout succeeds. Avg@8 reflects rollout reliability, while Pass@8 measures whether repeated sampling can recover a successful trajectory.

3pt

Table 1 shows that TextCraft mainly stresses multi-step subgoal maintenance. The base policy already solves many depth-1 tasks, but performance drops sharply as recipe depth increases. SFT greatly improves both metrics, and RL from the SFT model is consistently stronger than RL from the base policy. Among prefix-based methods, replay alone is weaker than the strongest SFT-initialized RL baselines, whereas Prefix-GRPO achieves the best overall Pass@8 and matches the strongest depth-4 result. This suggests that explicit prefix optimization is most useful for recovering successful continuations from difficult intermediate states, while strong SFT-initialized RL remains competitive in average rollout quality.

2pt

Table ¿tbl:tab:babyai95tasks? shows a shorter-horizon instruction-following regime where many methods already achieve high best-of-eight success. The base policy obtains strong Pass@8 on GoTo and AOD, but its Avg@8 remains weak on Pickup, Find Room, and SLoc, indicating unstable single-rollout execution. Distillation and RL both improve this average behavior: DAPO and SFT\(\rightarrow\)GRPO reach the strongest overall Avg@8, while Prefix-GRPO attains the best overall Pass@8. Compared with Replay-GRPO, explicit prefix optimization raises overall Avg@8 from 92.92 to 94.31 and Pass@8 from 97.78 to 100.00, with the clearest gains on Find Room and SLoc. Thus, BabyAI does not show the same large separation as TextCraft, but it still supports the reliability benefit of prefix optimization under multi-sample evaluation.

2pt

Table ¿tbl:tab:alfworld95tasks? shows that ALFWorld is substantially harder than BabyAI. The base policy obtains non-trivial Pass@8 on Simple Place and Examine, but average success remains low on multi-step manipulation families. SFT improves both overall metrics, and RL from the base model gives a much larger gain: GRPO, GRPO-MIS, and DAPO all improve overall Pass@8 from 53.00 to roughly 87–92. SFT-initialized RL further improves average rollout reliability, with SFT\(\rightarrow\)GRPO-MIS achieving the best overall Avg@8. Prefix-GRPO, however, obtains the best overall Pass@8 and substantially improves over Replay-GRPO, raising Avg@8 from 52.62 to 70.31 and Pass@8 from 80.00 to 94.50. The gains are most visible in best-of-eight recovery on Two-Obj Place and Heat Place. This pattern is consistent with TextCraft and BabyAI: explicit prefix optimization improves recovery from teacher-derived intermediate states, while the strongest SFT-initialized RL baselines can remain more reliable on average rollout success.

Overall, the task-family results show that the three environments stress different failure modes: symbolic subgoal maintenance in TextCraft, grounded instruction execution in BabyAI, and long-horizon household manipulation in ALFWorld. Across all three environments, Prefix-GRPO achieves the strongest overall Pass@8, but it does not always achieve the strongest overall Avg@8. This separates two effects: explicit prefix optimization improves multi-sample recovery from useful teacher-derived states, while strong SFT-initialized RL baselines can remain more reliable on average rollout success. The consistent gap between Replay-GRPO and Prefix-GRPO further indicates that teacher prefixes are more useful when they are directly optimized rather than only replayed as context.

4.3 Ablation Studies↩︎

4.3.0.1 Prefix construction.

The effectiveness of Prefix-GRPO depends on which intermediate states are exposed to the student. Table 2 compares four construction families on TextCraft. Fixed-ratio strategies split each teacher trajectory at manually chosen relative depths, while entropy-based strategies select cut points according to uncertainty patterns under the policy-distilled SFT model. The proposed Entropy-Change Top-3 construction performs best across all reported TextCraft metrics. Its advantage over Raw-Entropy Top-3 suggests that the most useful cut points are not necessarily the positions with the highest uncertainty, but the positions where the model’s uncertainty changes sharply after smoothing.

7pt

4.3.0.2 Replay versus explicit prefix optimization.

This ablation separates the effect of replaying teacher prefixes from the effect of applying policy-gradient updates to the prefix tokens themselves. Table 3 compares a replay-only variant with the full Prefix-GRPO objective under the same TextCraft data, rollout budget, optimization budget, and evaluation protocol. In the replay-only variant, teacher prefixes are used to reach recovered intermediate states and are included in the prompt context, but the policy loss is applied only to continuation tokens. Replay-only training exposes the policy to teacher-recovered intermediate states, but it remains substantially below full Prefix-GRPO; the largest gap appears at Pass@8, where explicit prefix optimization improves performance from 78.79 to 96.00. These results indicate that teacher prefixes should not only be replayed as context, but also receive direct policy updates.

7pt

4.3.0.3 Prefix-side signal design.

Table 4 studies two related design choices for using teacher prefixes:

• Prefix advantage assignment keeps the same group-query data but changes how prefix-token updates are weighted. Besides constant weighting and positive-only continuation signals, we test success-gated updates, family lift over the raw query, and the main continuation-magnitude anchor.

• Full-demo query reuse adds the complete teacher trajectory as an additional group-query member, with SFT-computed old log-probabilities and \(\mathrm{clamp}(R_{\mathrm{demo}}-\bar{R}_{\mathrm{raw}},0,1)\) as its demo advantage.

Here success-gated updates apply prefix loss only when the continuation succeeds, while family lift compares each prefix query against the raw no-prefix query from the same teacher trajectory. Full definitions are given in Appendix 8.1.0.1. The full-demo query improves low-budget metrics, but the continuation-magnitude anchor gives the best Pass@4 and Pass@8, suggesting that prefix updates are most useful when their strength is tied to the downstream continuation signal.

5pt

4.3.0.4 Prefix objective design.

Beyond the choice of cut states and prefix-side signals, Prefix-GRPO also depends on how prefix and continuation losses are combined. Table 5 compares the main split objective with a joint-objective variant and a weaker prefix-loss weight. The split objective with unit prefix weight performs best across all reported metrics, while the joint variant is substantially weaker. This suggests that prefix and continuation tokens benefit from sharing the same clipped-ratio form, but should retain separate loss terms and weighting. Additional prefix-clip sensitivity is reported in Appendix 8.1.

6pt

4.3.0.5 Training-side prefix diagnostics.

The previous ablations evaluate final task success. To verify that the prefix branch is also active during optimization, Table 6 reports prefix-side diagnostics averaged over the last 100 optimizer steps. Across construction variants, the prefix KL remains non-zero while the prefix clip fraction stays low, indicating that assistant prefix tokens receive controlled updates under the SFT old-policy anchor rather than being either ignored or destabilized. The loss share and token fraction further show that the prefix branch contributes a small but persistent part of the update, instead of dominating the continuation objective. Entropy-Change Top-3 also combines a moderate prefix loss share with the strongest Pass@8, suggesting that the best construction rule is not simply the one with the largest prefix update, but the one that yields useful and stable prefix-side learning signals.

6pt

5 Conclusion and Future Work↩︎

In this paper, we introduced Prefix-GRPO, a reinforcement learning framework that turns teacher-trajectory prefixes into explicit optimization targets rather than treating them as passive context. Our method combines group-query prefix construction with prompt-space prefix old-logprob anchoring, allowing a single teacher trajectory to produce multiple replay-aligned training states and enabling historical prefix tokens to participate directly in GRPO optimization. Across TextCraft, BabyAI, and ALFWorld, the experiments show that small-model agents benefit from combining distillation with reinforcement learning, while the TextCraft ablations provide the clearest evidence for the role of prefix optimization: replaying teacher prefixes alone is not sufficient, and directly optimizing prefix tokens improves multi-sample recovery on harder tasks.

Several limitations remain: most results are single-run experiments, Prefix-GRPO requires replay-verifiable teacher prefixes, and the SFT old-policy anchor still needs more systematic study. Future work should extend matched ablations across more environments and analyze alternative prefix-side objectives and old-policy anchors.

A Datasets to 0.6em.

A.1 Dataset Overview to 0.6em.

A.2 Prefix-GRPO Data Construction to 0.6em.

A.3 Mask Alignment Example to 0.6em.

B Implementation Details in Each Environment to 0.6em.

B.1 TextCraft to 0.6em.

B.2 BabyAI to 0.6em.

B.3 ALFWorld to 0.6em.

B.4 Prompt Templates to 0.6em.

C Additional Experimental Results to 0.6em.

C.1 Additional TextCraft Objective Ablations to 0.6em.

C.2 BabyAI Construction Ablation to 0.6em.

D Behavior Diagnostics and Case Studies to 0.6em.

D.1 Training Dynamics to 0.6em.

D.2 Evaluation Diagnostics to 0.6em.

D.3 Failure-Mode Diagnostics to 0.6em.

D.4 Rollout Visualization to 0.6em.

6 Datasets↩︎

6.1 Dataset Overview↩︎

Table ¿tbl:tab:dataset95statistics? summarizes the task-family splits used for GRPO training, prefix fullflow training, and evaluation in the three retained environments. We then describe how teacher trajectories are converted into replay-validated Prefix-GRPO training samples.

2pt

6.2 Prefix-GRPO Data Construction↩︎

Prefix-GRPO first turns teacher interaction trajectories into replay-validated training states before online reinforcement learning begins. The data construction pipeline consists of five stages: collecting teacher trajectories, constructing group-query prefixes, validating each prefix by environment replay, canonicalizing the prompt at the cut state, and estimating old log-probabilities for historical prefix tokens. The resulting corpus stores both the textual prompt and the sidecar fields needed by the Prefix-GRPO objective.

6.2.0.1 Group-query construction.

Let \((x_1,\ldots,x_T)\) denote the token sequence of a fully serialized teacher trajectory. We use the same policy-distilled SFT checkpoint that later serves as the prefix old policy to compute teacher-forcing next-token entropy: \[\label{eq:entropy} H_t=-\sum_v p_\phi(v\mid x_{<t})\log p_\phi(v\mid x_{<t}).\tag{7}\] In the main TextCraft pipeline, scoring is restricted to the interaction_assistant token domain. After filtering to this domain and sorting by original token position, we smooth the entropy sequence with a centered moving average. The main experiments use \(w=11\) as the main instantiation rather than as a sensitivity-tuned optimum. Let \(\hat{H}_j\) denote the smoothed entropy on the filtered sequence. The entropy-change score is then \[\label{eq:entropy-change} s_j=\left|\hat{H}_j-\hat{H}_{j-1}\right|.\tag{8}\] Candidate tokens are ranked globally by \(s_j\), mapped to assistant turns, deduplicated within each assistant turn, and truncated to at most top-\(k=3\) distinct assistant turns. If multiple high-scoring tokens map to the same assistant turn, only the highest-scoring token is kept for that turn. We impose no minimum temporal gap across selected turns; if fewer than three distinct assistant turns are available, the trajectory contributes fewer prefix queries. In addition to these prefix queries, we keep one raw no-prefix query for the same teacher trajectory. We use \(k=3\) as the main instantiation to expose multiple cut states while keeping dataset expansion manageable, rather than claiming it as a sensitivity-tuned optimum.

6.2.0.2 Replay validation.

Replay validation is stricter than checking whether prefix actions are merely executable. For each candidate prefix query, we create a TextCraft environment using the sample’s task identifier and goal, replay the teacher prefix actions in order, and compare the replayed observation against the first continuation-side user observation. If the continuation still contains another assistant action, we additionally execute that action as a one-step probe and compare the next observation as well. Matching is performed on shared structured fields extracted from the observations, including Inventory:, Got ..., and Crafted ....

Table 7 compares replay validation quality with downstream TextCraft performance across construction families. Replay validation provides a data-quality floor, but the downstream result is not monotonic in the validated-prefix ratio: Raw-Entropy Top-3 has the highest validation rate, whereas Entropy-Change Top-3 gives the strongest Pass@8.

2pt

6.2.0.3 Canonicalized prompt format.

Suppose a selected cut point falls after assistant turn \(k\). Then prefix_messages contain all messages up to and including that assistant message, while continuation_messages begin from the following user observation. After replay validation, the training prompt is canonicalized as \[\label{eq:canonical-prompt-app} x_i=\mathrm{concat}(\text{system},\text{prefix history},\text{cut-state user observation}).\tag{9}\] The prompt therefore ends at the replayed user-side observation rather than at the teacher action itself.

6.2.0.4 Prefix old-logprob estimation.

Continuation tokens are produced by the online rollout and therefore naturally come with old log-probabilities from the rollout policy. Prefix tokens are different: they are teacher-derived assistant tokens already located inside the prompt history, not tokens sampled by the current online rollout. To support PPO/GRPO-style optimization on these tokens, we use a policy-distilled SFT checkpoint as the prefix old policy. For each canonicalized prompt \(x_i\), we run teacher forcing under the SFT old policy and save the prefix sidecar \[\label{eq:prefix-sidecar-app} \mathcal{S}_i=(b_i,e_i,m_i,\ell_i^{\mathrm{old}},n_i),\tag{10}\] where \([b_i,e_i)\) is the prefix span, \(m_i\) is the prefix mask, \(\ell_i^{\mathrm{old}}\) are cached old log-probabilities stored densely over the span, and \(n_i=\sum_u m_i[u-b_i]\) is the prefix token count. The valid prefix-token positions are \[\label{eq:prefix-mask-set-app} M_i^{\mathrm{pre}}=\{u\mid b_i\le u<e_i,\; m_i[u-b_i]=1\}.\tag{11}\] We additionally verify that prefix_token_count equals prefix_mask.sum() and that both prefix_mask and assistant_prefix_old_log_probs have length \(e_i-b_i\).

6.3 Mask Alignment Example↩︎

Tables 8 and 9 give a concrete replay-validated TextCraft example. The key point is that span, mask, and old log-probabilities are all generated after prompt canonicalization and therefore share the same prompt-token coordinate system.

4pt

3pt

This example shows that changes in system or observation tokens do not require reusing stale turn-level offsets: the sidecar is regenerated on the final canonicalized prompt, and prefix loss is applied only to mask-one assistant positions.

7 Implementation Details in Each Environment↩︎

We describe the environment-specific task setup, evaluation protocol, and training configuration below.

7.0.0.1 Compute resources.

All main experiments were run on 2 NVIDIA H200 GPUs. Under this setup, an SFT run typically took about 1–3 hours, while an RL training run took approximately 1–2 days depending on the environment and ablation setting. Evaluation used the corresponding trained checkpoint with a fixed eight-rollout sampling budget per task.

7.1 TextCraft↩︎

7.1.0.1 Task setup and evaluation.

TextCraft is the primary environment used to validate the full Prefix-GRPO pipeline. The agent observes a symbolic crafting task and interacts through textual actions such as inventory inspection, acquiring primitive items, and crafting target objects from recipes. A trajectory is successful only when the target object is actually produced in the environment; merely emitting a completion phrase is not sufficient.

7.1.0.2 Prompt format.

TextCraft uses a fixed system prompt followed by a dynamic task observation. Prefix rows additionally include the canonicalized prefix history up to the cut state. The prompt template is listed in Appendix 7.4.

7.1.0.3 Training configuration.

TextCraft uses the full Prefix-GRPO pipeline, including replay validation, prefix old-logprob recomputation, and online rollout from recovered cut states. Table 10 lists the main RL training settings; the SFT stage uses 20 epochs and learning rate \(1\times 10^{-5}\).

6pt

7.1.0.4 Runtime and infrastructure.

TextCraft training requires an environment server, a rollout backend, replay validation, and old-logprob recomputation. The main additional preprocessing cost of Prefix-GRPO comes from replay validation and prompt-space teacher forcing under the SFT old policy. These costs are incurred before RL training and are amortized across subsequent policy updates.

7.2 BabyAI↩︎

7.2.0.1 Task setup and evaluation.

BabyAI evaluates grounded instruction following in grid-world tasks. The agent receives a natural-language instruction and must execute a sequence of actions that satisfies the instruction in the environment. The main paper reports Avg@8 and Pass@8 by official task family and also keeps the overall success rate for supplementary comparison.

7.2.0.2 Prompt format.

BabyAI uses a thought-action-observation prompt with an explicit action format and the currently available action set. The prompt template is listed in Appendix 7.4.

7.2.0.3 Training configuration.

Table 11 lists the BabyAI RL settings, with rollout limits adapted to the grid-world interaction format. The SFT stage uses 20 epochs and learning rate \(1\times 10^{-5}\).

6pt

7.3 ALFWorld↩︎

7.3.0.1 Task setup and evaluation.

ALFWorld evaluates text-based household instruction following with grounded task structure. The agent must interpret a household goal, issue valid textual actions, and reach the required terminal state. In the main paper, results are reported by official task family using Avg@8 and Pass@8, which exposes whether gains are concentrated on short placement tasks or extend to harder multi-step tasks.

7.3.0.2 Prompt format.

ALFWorld uses a household-task system prompt, a fixed user message, and environment observations from reset and step outputs. The prompt template is listed in Appendix 7.4.

7.3.0.3 Training configuration.

Table 12 lists the ALFWorld RL settings, including the longer context budget used for household trajectories. The SFT stage uses 20 epochs and learning rate \(1\times 10^{-5}\).

6pt

7.4 Prompt Templates↩︎

We collect the full prompt templates in one place to avoid interrupting the environment-specific configuration sections with long verbatim blocks. The serialized training prompts additionally include dataset-specific prefix histories and cut-state observations.

TextCraft.

You are a Minecraft Assistant. Your goal is to craft items by managing resources and recipes.

CORE PROTOCOL:
1. THINK FIRST: before any action, analyze the current state.
2. ONE ACTION: output exactly one action per turn.
3. BOX FORMAT: wrap your command in [[ ]], e.g., Action: [[ inventory ]]
4. NO HALLUCINATION: do not simulate the environment response.

REASONING LOGIC:
When trying to acquire an item [Target]:
1. Check inventory. If you already have [Target], proceed.
2. Check recipe. If [Target] has a recipe, get missing ingredients first.
3. If [Target] is a base material, use get [Target].
4. If get or craft fails, check recipes and exact item names again.

CORE COMMAND SET:
- craft [target] using [ingredients]
- get [item]
- inventory

INTERACTION EXAMPLE:
[Environment]
Goal: craft powered rail
Recipes:
- craft 6 powered rail using 1 redstone, 1 stick, 6 gold ingot
- craft 1 gold ingot using 9 gold nugget
- craft 1 stick using 2 bamboo

[You]
Think: Recipe for powered rail needs redstone, stick, gold ingot.
I need to check what I have.
Action: [[ inventory ]]

[Environment]
Inventory: [redstone] (1)

[You]
Think: I have redstone. I miss stick and gold ingot. Gold ingot
needs gold nugget, so I should get gold nugget first.
Action: [[ get 54 gold nugget ]]

[Environment]
Got 54 gold nugget.

[You]
Think: Now I have nuggets. I can craft gold ingots.
Action: [[ craft 6 gold ingot using 54 gold nugget ]]

[Environment]
Crafted 6 gold ingot.

[You]
Think: Now I need a stick. Recipe says stick needs bamboo.
Action: [[ get 2 bamboo ]]

[Environment]
Got 2 bamboo.

[You]
Think: I have bamboo. I can craft the stick.
Action: [[ craft 1 stick using 2 bamboo ]]

[Environment]
Crafted 1 stick.

[You]
Think: I have redstone, gold ingots, and stick. I can finally craft
the powered rail.
Action: [[ craft 6 powered rail using 1 redstone, 1 stick, 6 gold ingot ]]

[Environment]
Crafted 6 powered rail.

[You]
Think: Goal achieved.
Action: [[ Task Completed! ]]

Dynamic user observation:
Crafting commands:
{crafting_commands}

Goal: craft {goal_item}.

BabyAI.

You are an exploration master in the BabyAI environment. Every round I
will give you an observation and the currently available actions. Your
goal is to complete the given navigation or object-manipulation task.

Common actions include:
- turn left
- turn right
- move forward
- go to <object> <id>
- pickup <object> <id>
- drop
- toggle
- go through <door> <id>
- toggle and go through <door> <id>
- check available actions

Your response should use exactly this format:
Thought:
your thoughts.

Action:
[[ your next action ]]

IMPORTANT: Output exactly one executable action and always wrap it in [[ ]].

Dynamic user observation:
Your goal: {mission}
{natural_language_grid_observation}
Available actions: [{available_actions}]

ALFWorld.

System prompt:
You are an agent in the ALFWorld environment. Your goal is to complete
household tasks by interacting with objects in various rooms.

You can use actions like:
- go to [location]
- take [object] from [location]
- put [object] in/on [location]
- open [object]
- close [object]
- toggle [object]
- clean [object] with [location]
- heat [object] with [location]
- cool [object] with [location]
- examine [object]
- inventory

Instructions:
- Read the task instruction carefully.
- Plan your actions to complete the task.
- Navigate between rooms as needed.
- Manipulate objects correctly.

Your response should contain only the action.

Fixed user message:
Please help me complete the household task. I will provide you with the task description.

Dynamic environment observation:
{room_observation}
Your task is to: {household_task}

8 Additional Experimental Results↩︎

8.1 Additional TextCraft Objective Ablations↩︎

The main paper reports the ablations that directly test trajectory construction, replay, prefix-signal assignment, loss combination, and prefix-side training diagnostics. Here we provide lower-level TextCraft sensitivity studies for prefix clipping.

8.1.0.1 Definitions of prefix-side signals.

For completeness, we define the prefix-side signals used in Table 4:

• Constant assigns every valid prefix token a unit positive weight.

• Positive continuation clips the continuation-derived prefix weight at zero.

• Success-gated applies a unit prefix weight only when the continuation rollout succeeds, and sets the prefix weight to zero otherwise.

• Family lift groups queries generated from the same teacher trajectory into one raw no-prefix query and several prefix queries; it computes the lift of a prefix query relative to the raw query and combines this non-negative lift with the continuation-magnitude signal.

• Continuation-magnitude anchor uses the absolute continuation signal as the prefix weight.

• Full-demo query adds the complete teacher trajectory as an additional query in the same group, uses SFT-computed old log-probabilities for its assistant tokens, and assigns it the demo advantage \(\mathrm{clamp}(R_{\mathrm{demo}}-\bar{R}_{\mathrm{raw}},0,1)\).

8.1.0.2 Prefix clip settings.

Table 13 summarizes the completed prefix-clip sweep. The main symmetric range \(0.2/0.2\) gives the strongest Pass@8; tighter, looser, or asymmetric settings generally reduce multi-sample recovery. The prefix-only \(0.28/0.2\) and asymmetric \(0.5/0.2\) settings remain competitive but do not match the main setting. These results indicate that prefix updates benefit from clipping, but overly permissive clipping can weaken the old-logprob anchor.

6pt

8.2 BabyAI Construction Ablation↩︎

8.2.0.1 Prefix construction.

Table 14 reports the BabyAI construction-family ablation together with replay-validation statistics. All four construction families achieve similarly high replay-validation rates, ranging from 97.68 to 98.07, which suggests that replayability is not the main bottleneck in this shorter-horizon environment. However, downstream RL performance still differs across construction rules: Entropy-Change Top-3 obtains the strongest results across all reported metrics, reaching 100.00 Pass@8. This mirrors the TextCraft finding that construction quality is not determined solely by whether a prefix can be replayed, but also by whether the selected cut states provide useful learning signals for subsequent RL.

2pt

9 Behavior Diagnostics and Case Studies↩︎

9.1 Training Dynamics↩︎

Figure 3 visualizes the TextCraft main-change-top3 training run. The first 20 steps are dominated by long responses and many interaction turns, while the last 20 steps are much shorter and more successful: mean reward increases from 0.489 to 0.835, mean turns decrease from 26.60 to 11.06, and mean response length decreases from 2863 to 426 tokens. This suggests that Prefix-GRPO changes both task success and interaction style.

9.2 Evaluation Diagnostics↩︎

The final checkpoint is evaluated on 800 samples covering 100 TextCraft tasks, with no missing samples, no generation or environment errors, and no goal mismatches. Figure 4 (a) shows the depth-wise breakdown: Prefix-GRPO solves nearly all depth-1 and depth-2 tasks under Pass@8, while the remaining errors concentrate in deeper recipes. Depth 3 has 0.50 Avg@8 and 0.92 Pass@8; depth 4 remains difficult with 0.042 Avg@8 and 0.333 Pass@8.

9.3 Failure Modes↩︎

Among 800 evaluated samples, 653 succeed and 147 fail. Successful rollouts require 8.02 turns on average, while every failed rollout reaches the 30-turn limit. Figure 4 (b) summarizes final-observation categories for failed trajectories. Most failures are not parser failures: only one failed sample ends with an invalid action. The dominant failures are resource-state errors, including unavailable items (46 cases), insufficient inventory for craft (42 cases), crafted non-terminal items (30 cases), and unfinished item acquisition (26 cases).

9.4 Rollout Visualization↩︎

Figure 5 shows representative TextCraft rollouts from the trained Prefix-GRPO policy. The successful composter trajectory illustrates concise state tracking: after acquiring the required slabs, the policy immediately executes the terminal craft action. The failed magma block trajectory illustrates a remaining failure mode: the policy makes partial progress but then enters a repeated ingredient-recovery loop and reaches the turn limit.

References↩︎