The policy-gradient loss

The policy-gradient loss is the token-level quantity a trainer differentiates to apply a policy-gradient update. It weights the log-probability of each token the policy sampled by that token's advantage, restricted by a loss mask to tokens the policy authored and divided by a normalizer:

\mathcal{L}(\theta) = -\frac{1}{Z}\sum_{i}\sum_{t} m_{i,t}\,A_{i,t}\,\log\pi_\theta(y_{i,t}\mid x_i, y_{i,<t}).

Here $i$ indexes rollouts in the batch, $t$ indexes token positions, $m_{i,t}\in\{0,1\}$ is the mask, and $Z$ is the normalizer. In an agent trace many tokens come from the environment rather than the policy, so the mask and the normalizer determine what the update optimizes.

Loss masks and token authorship

An agent trace renders many sources into one token sequence: a system prompt, the task, model turns, tool calls, tool results, harness scaffolding, and sometimes a verifier's message. Only some positions were sampled from the policy being trained.

Span	Author	In context	In the loss ( $m=1$ )
System prompt and task	taskset / harness	yes	no
Model reasoning and reply	policy	yes	yes
Tool call arguments	policy	yes	yes
Tool output, file contents, test logs	environment	yes	no
Role tags and turn openers added by the template	renderer	yes	no

The policy-gradient term trains only sampled actions. Environment observations stay in the context the next action conditions on and are masked from that term, because the environment's transition probabilities do not depend on $\theta$ and contribute no gradient (policy gradients). Search-R1 (Jin et al.) applied this to retrieved passages: in their PPO ablation on Qwen2.5-7B-Base, masking retrieved tokens from the loss raised average exact match over seven question-answering datasets from 0.343 to 0.431. A separate auxiliary objective, such as ECHO, can still train the policy to predict observations.

Mask errors are silent and damaging. A missing mask attaches policy-gradient terms to tool output, error messages and user text, which were never sampled actions. A mask shifted by one position attaches each advantage to the wrong decision. A context re-tokenized differently from what was sampled evaluates every probability under the wrong prefix, which is why systems keep the sampled token IDs as the source of truth (token-level rendering).

Masks also drop tokens on purpose. Some recipes mask rollouts truncated at the length limit instead of scoring them as failures, since a cut-off answer is ambiguous evidence. Off-policy corrections mask tokens whose probability has drifted too far from the sampler's (importance sampling).

Token versus sequence normalization

The normalizer $Z$ sets how much each rollout and each token contributes. Three conventions are common:

Convention	Per-token weight in rollout $i$	Used by
Sequence mean, then batch mean	$\frac{1}{N\,L_i}$	original GRPO
Token mean over the batch	$\frac{1}{\sum_j L_j}$	DAPO
Constant	$\frac{1}{N\,L_{\max}}$	Dr. GRPO

where $L_i$ is the number of masked-in tokens in rollout $i$ , $N$ is the number of rollouts, and $L_{\max}$ is a fixed budget such as the generation limit.

Sequence-mean normalization gives every rollout equal total weight regardless of length. Token-mean normalization gives every token equal weight, so long rollouts contribute proportionally more. A constant normalizer also gives every token equal weight and keeps that weight independent of what else is in the batch.

Length bias

Take two rollouts in one group: a correct answer with 100 tokens ( $A = +1$ ) and an incorrect one with 1,000 tokens ( $A = -1$ ).

Convention	Weight per token, correct	Weight per token, incorrect
Sequence mean	$\frac{1}{2\cdot 100} = 0.005$	$\frac{1}{2\cdot 1000} = 0.0005$
Token mean	$\frac{1}{1100} \approx 0.00091$	$\frac{1}{1100} \approx 0.00091$

Under sequence-mean normalization each token of the long wrong answer is penalized ten times less than each token of the short right answer is rewarded. The pressure against a bad behavior thins out the longer the rollout runs, so incorrect responses drift longer and correct ones shorter. Liu et al. identify this response-level length bias in GRPO and in several open-source PPO implementations, and argue that the response-length growth often read as emergent long reasoning in R1-style training is partly confounded by it.

Token-level normalization removes that asymmetry but does not make length neutral. A long rollout has more tokens and so more total influence on the update. Whether that is right depends on what counts as an example: a twenty-turn agent rollout may contain more consequential decisions than a short answer, and it also contains more routine formatting. A tensor.mean(-1) in the loss code makes that choice implicitly.

Explicit length penalties in the reward are a separate lever. They change which rollouts rank above the baseline, while normalization only changes how hard each ranked rollout pushes.

Auxiliary terms

The reward-weighted term is often combined with others:

\mathcal{L} = \mathcal{L}_{\mathrm{PG}} + \beta\,\mathcal{L}_{\mathrm{KL}} - c_{H}\,\mathcal{H} + c_{\mathrm{CE}}\,\mathcal{L}_{\mathrm{CE}}.

A KL penalty keeps the policy near a reference or behavior distribution.
An entropy bonus $\mathcal H$ resists premature concentration of the policy (entropy collapse).
A cross-entropy term imitates chosen targets, such as demonstrations or, with ECHO, observations.

Each term has its own natural count of tokens. Summing them into one scalar with one shared denominator lets the component with the most tokens set the effective learning rate of the others. Normalizing each by its own token count and logging each separately makes it possible to tell which pressure caused a change.

Observation prediction (ECHO)

Masking observations is the default, and it discards information. Terminal output, error messages and test logs record how the environment responded to the agent's actions, and a failed rollout with zero advantage still contains all of it.

ECHO (Environment Cross-entropy Hybrid Objective, Shrivastava et al.) keeps the standard policy-gradient loss on action tokens and adds a weighted cross-entropy loss that trains the policy to predict selected observation tokens:

\mathcal{L}_{\mathrm{ECHO}} = \mathcal{L}_{\mathrm{GRPO}}(\text{actions}) \;-\; \frac{c_{\mathrm{CE}}}{|\mathcal{O}|}\sum_{t\,\in\,\mathcal{O}'} \log\pi_\theta(o_t \mid x, y_{<t}).

Here $\mathcal{O}$ is the set of all observation positions in the sequence and $\mathcal{O}' \subseteq \mathcal{O}$ is the subset selected as prediction targets. Dividing by $|\mathcal{O}|$ keeps the auxiliary term on a per-observation-token scale, so long and short observations get comparable weight under one coefficient.

The observation tokens are already in the forward pass, so the extra loss costs no additional rollouts. The authors report that it roughly doubles GRPO's pass@1 on TerminalBench-2.0, from 2.70% to 5.17% for Qwen3-8B and from 5.17% to 10.79% for Qwen3-14B. Without any demonstrations, it matches expert SFT followed by GRPO on the authors' held-out terminal tasks and closes about half of that gap on TerminalBench-2.0.

Observation prediction helps when predicting the observation teaches useful dynamics, such as what a shell command will print or which test will fail. It can waste capacity when observations are arbitrary retrieved text, where prediction amounts to memorization without improving the decision to search, read or stop. The choice is per token source: action, context, or prediction target, encoded in the mask and the loss weights.

In prime-rl

The prime-rl trainer receives masked samples that the orchestrator converts from traces (training systems) and computes its loss as three components, rl, ce and ref_kl, each summed over its own tokens and divided by that component's global token count across ranks (loss.py). The rl term is therefore token-mean normalized over the step, and zero-advantage tokens are removed from it, denominator included.

ECHO is an algorithm type that runs GRPO on action tokens and routes observation tokens into ce:

[orchestrator.algo]
type = "echo"

[orchestrator.algo.roles.tool]
alpha = 0.1

Because ce has its own denominator, adding observation tokens leaves the per-token scale of the RL term unchanged. That denominator is the step's global count of selected ce tokens, which differs from the paper's per-sequence division by total observation length.