Curated Agent Traces with Tool-Response Masking

Why People Wanted It

A model trained only on natural software artifacts can learn code, but those artifacts do not show the rhythm of interacting with tools: make a call, inspect the result, test a change, recover from failure. Agent traces can supply that missing behavior. But a transcript contains both the assistant’s decisions and outputs produced by the environment. If the model is trained to reproduce both indiscriminately, it is being asked to imitate things it did not choose or generate.

So the real question was not simply whether to add agent traces. It was how to expose the model to interaction patterns without making tool output part of the target it must imitate—and how to keep traces useful without allowing them to displace natural data.

The Core Move

Use a bounded share of diverse agent trajectories in midtraining, and calculate learning loss on assistant turns—including tool calls—but not on tool responses.

The stated recipe gives no formula for the trace share or the loss mask. In words, the mask says: learn from what the assistant says and does; use the tool’s response as context for the next assistant action, not as text the model is trained to produce.

Why a Bounded Share Keeps Traces in Their Role

Here is the first deeper layer: the recipe treats agent traces as a supplement to a larger natural-data blend, not as a replacement for it. The motivation is practical. Natural software artifacts teach code and repository structure; trajectories add interaction patterns that those artifacts lack. The reported blend-validation experiment compares a full blend against a version without agentic traces, and describes natural data as the large majority of the blend.

That balance matters because traces are not interchangeable with ordinary code examples: they carry scaffold instructions, tool definitions, and environment-specific conventions. The corpus spans more than ten scaffolds, and the text notes that traces from one scaffold transfer poorly to others. Diversity is therefore not decoration; it gives the model exposure to different interaction formats. The paper does not state a universal numerical share for agent traces here, so the transferable point is the bounded, diverse role—not a fixed percentage.

Why Tool Responses Are Context, Not Targets

Now the technical distinction: a trajectory alternates between assistant turns and tool responses. Training only on assistant turns means the tool output still appears in the conversational context, where it can inform what the assistant does next, but it is excluded from the prediction loss. Assistant tool calls remain targets: choosing and formatting an action is part of the behavior being learned.

This is a selective target, not a selective input. For example, after a test command returns an error, the model can learn to respond with a repair or another diagnostic call without being trained to emit the test runner’s error text itself. The stated rationale is empirical: also training on tool responses hurts results after SFT. That finding supports the masking choice; it does not imply that tool outputs should be removed from the trajectory.

Minimal Causal Code

The loop below shows the causal distinction: every turn is available as context, while only assistant turns contribute to the loss. A production pipeline would need to map the dataset’s role labels and token boundaries to this simple mask.

import torch
import torch.nn.functional as F

for batch in loader:
    input_ids = batch["input_ids"].to(device)
    assistant_mask = batch["assistant_mask"].to(device)  # 1 for assistant tokens, 0 for tool responses

    logits = model(input_ids[:, :-1]).logits
    targets = input_ids[:, 1:]
    loss_mask = assistant_mask[:, 1:]

    token_loss = F.cross_entropy(
        logits.reshape(-1, logits.size(-1)),
        targets.reshape(-1),
        reduction="none",
    ).view_as(targets)

    loss = (token_loss * loss_mask).sum() / loss_mask.sum().clamp_min(1)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

What You Gain, and What You Lose

The gain is a training target aligned with the model’s role: assistant decisions and tool calls are learned, while tool responses remain evidence the assistant can condition on. A bounded, varied trace slice can add multi-step interaction patterns that natural software artifacts do not contain.

The cost is that the model is not directly optimized to reproduce tool responses, and trace diversity does not make every scaffold convention universal. The reported result that training on tool responses hurts after SFT is a reason to prefer masking in this recipe, not a guarantee that the same choice is optimal for every data or training setup. Trace curation also matters: the source describes filtering shortcuts and undesirable patterns such as malformed calls and repetitive inspection loops. Natural software-engineering data and repository construction are covered in separate modules; here the boundary is the trace’s share and its loss targets.