Module 6. The LLM as a policy: tokens, log-probabilities, and KL

Stop reason and task semantics

The contract of this particular task says that a model-selected EOS is a natural terminal event, whereas reaching max_new_tokens is an external truncation. The system returns two responses:

the model selects EOS on the 47th newly generated token;
generation stops on token 64 because max_new_tokens=64, without selecting EOS.

For each response, specify stop_reason, terminated, and truncated. Then explain:

why another task could map the same raw reasons differently;
why response length does not recover the stopping cause;
why stop_reason must be stored separately from the two Boolean flags.

Show answer and solution

Under the stated task contract:

stop_reason="eos", terminated=True, truncated=False;
stop_reason="length", terminated=False, truncated=True.

Those flags follow from the task contract, not from the raw reason name alone. A stop string can denote a genuine task completion in one system and merely a technical delimiter in another. A true terminal event may also coincide with a rollout-buffer boundary, so a generic data contract should not automatically prohibit both flags from being true.

The same returned length can result from EOS, a token budget, a timeout, or an external interrupt. Store the raw machine-readable reason, then map it explicitly to the task-specific termination and truncation semantics.

Is the prefix a sufficient state?

Consider two processes.

An LLM generates plain text without external calls.
An LLM emits a web-search call, and the environment returns a result that depends on the current network state.

In which process is $s_t=(x,y_{<t})$ sufficient for a Markov model? What must be added in the second process? Separate model actions from environment observations.

Show answer and solution

For plain generation, $(x,y_{<t})$ is sufficient when the prompt template and generation constraints are fixed. The transition after a token is deterministic concatenation.

The tool-using process also needs tool output and relevant external state. The tool call is a model action; the returned document is an environment observation. It is not a policy-selected token merely because the interface serializes everything into one transcript.

The one-token error

You have logits with shape $[B,L,V]$ and input_ids with shape $[B,L]$ . A developer computes

Python

selected = log_softmax(logits).gather(-1, input_ids[..., None])

without shifting. Why do the dimensions look plausible while the result is semantically wrong? State the correct alignment between a logits position and its target token.

Show answer and solution

Logits at position $i$ predict token $i+1$ , not token $i$ . Without shifting, the code pairs logits with the already-observed input at the same position. The correct alignment is

\text{logits}[:,i,:]\leftrightarrow\text{input\_ids}[:,i+1].

Use logits[:, :-1, :] and input_ids[:, 1:].

What does PPL say?

Model A has lower response PPL than model B. Can we conclude that:

A answers correctly more often;
A is more useful to a user;
A predicts the selected tokens better under the same tokenizer, template, and mask?

Answer each statement and give at least two reasons why PPL values from different experimental protocols may not be comparable.

Show answer and solution

More correct answers: no.
Greater user utility: no.
Better prediction of selected tokens under the same protocol: yes; lower PPL means lower mean negative log-probability.

Different tokenizers, masks, chat templates, datasets, or logarithm bases make comparisons ambiguous. A likely continuation can also be factually wrong.

“Negative KL”

For one sampled token,

\log\pi(a\mid s)-\log\pi_{\mathrm{ref}}(a\mid s)<0.

Does this mean KL between the policies is negative? Explain the difference between exact conditional KL and one sampled log-ratio.

Show answer and solution

No. Exact conditional KL is an expectation and is non-negative:

D_{\mathrm{KL}}(\pi\|\pi_{\mathrm{ref}}) = \mathbb{E}_{a\sim\pi} [\log\pi(a)-\log\pi_{\mathrm{ref}}(a)]\ge0.

One sampled log-ratio can be negative when the reference assigns more probability to that action. Only an average under the correct sampling distribution estimates KL.

Support conditions for k3

Let $x\sim q$ , $r(x)=p(x)/q(x)$ , and

k_3=(r-1)-\log r.

Separate two requirements:

when is $\log r$ finite on samples from $q$ ?
when does $\mathbb E_q[p/q]=1$ hold?

Why does the finite unbiased identity in this form require compatible supports? Give one counterexample in each direction.

Show answer and solution

For $\log r$ to be finite wherever $q$ can sample, we need

q(x)>0\Rightarrow p(x)>0,

or $q\ll p$ . If $q(x)>0$ but $p(x)=0$ , then $\log r=-\infty$ and $D_{\mathrm{KL}}(q\|p)=\infty$ .

The control-variate step uses

\mathbb E_q\!\left[\frac pq\right] =\sum_{x:q(x)>0}p(x).

This equals one only when $p$ has no mass outside the support of $q$ , namely $p\ll q$ . Thus a finite exact identity in the displayed discrete form is naturally stated for common positive support.

Counterexamples:

$q=(1,0)$ and $p=(0.8,0.2)$ give $\mathbb E_q[p/q]=0.8$ because the second mass of $p$ is never observed;
$q=(0.8,0.2)$ and $p=(1,0)$ make the log-ratio at the second point negative infinity, so the KL is infinite.

Numerical smoothing changes the estimand; it does not recover missing support for free.

Which policy generated the token?

A transformer emits base logits. The engine then applies temperature scaling, bans several tokens, applies top-p, and samples an action. The trainer stores a log-probability from the unmodified base softmax.

What error does this create in a PPO probability ratio? What metadata is needed to reconstruct the actual behavior policy?

Show answer and solution

The PPO denominator must be the action probability under the post-temperature, post-mask, post-top-p behavior policy. Base softmax has different probabilities and, after truncation, a different support. The ratio is systematically wrong.

Record the weight version, tokenizer version and template, temperature, all masks and penalties, top-k/top-p/min-p, processor order, stopping rules, seed, and the actual post-transform old_logprob.

A minimal rollout contract

You are designing a batch for reinforcement learning with verifiable rewards (RLVR). Select the fields required to reproduce a training objective reliably:

response text;
response_ids;
response_mask;
old_logprobs;
policy_version;
sampling_config;
reward_raw;
terminated, truncated, stop_reason;
the user-interface color.

Explain each selected field. Why is response text alone insufficient?

Show answer and solution

The required fields are response_ids, response_mask, old_logprobs, policy_version, sampling_config, reward_raw, and the three stop fields. Response text is useful for audit but does not replace tokens.

Token identifiers fix the actions; the mask separates actions from prompt and padding; old log-probabilities define the ratio denominator; the policy version and generation settings identify the behavior policy; raw reward is separated from regularizers; stop metadata defines the episode boundary. UI color is irrelevant. Retokenizing text under another template or tokenizer version can produce different actions.

Sequence log-probability and PPL

Three active response tokens had probabilities

0.5,\qquad 0.25,\qquad 0.8.

Compute:

sequence log-probability;
mean negative log-probability;
PPL.

What happens to the sum and PPL if a fourth token with probability $1$ is added?

Show answer and solution

L_{\mathrm{seq}} =\log0.5+\log0.25+\log0.8 =\log0.1 \approx-2.302585.

-\bar\ell=-L_{\mathrm{seq}}/3\approx0.767528,

so

\operatorname{PPL}=\exp(0.767528)\approx2.154435.

A token with probability one adds zero to the sum. The denominator becomes four, so PPL falls to $\exp(2.302585/4)\approx1.7783$ . This illustrates mask dependence.

The KL chain rule on a depth-two tree

At the first step,

p_0=(0.6,0.4),\qquad q_0=(0.5,0.5).

Second-step conditionals are

p(\cdot\mid0)=(0.7,0.3),\quad q(\cdot\mid0)=(0.4,0.6),

p(\cdot\mid1)=(0.2,0.8),\quad q(\cdot\mid1)=(0.5,0.5).

Compute

D_{\mathrm{KL}}(p_0\|q_0) + \mathbb{E}_{a_1\sim p_0} D_{\mathrm{KL}} \bigl(p(\cdot\mid a_1)\|q(\cdot\mid a_1)\bigr)

and verify that it equals KL between the two joint action distributions.

Show answer and solution

The root contribution is

D_{\mathrm{KL}}(p_0\|q_0) \approx0.0201355.

The expected conditional contribution is

0.6D_{\mathrm{KL}}((0.7,0.3)\|(0.4,0.6)) + 0.4D_{\mathrm{KL}}((0.2,0.8)\|(0.5,0.5)) \approx0.1873700.

Therefore

D_{\mathrm{KL}}(p_{1:2}\|q_{1:2})\approx0.2075056.

Constructing the four joint probabilities directly gives the same value.

k1, k2, and k3

For $r=p(x)/q(x)$ , compute k1, k2, and k3 at:

$r=2$ ;
$r=1/2$ .

Which estimator changes sign? Which two remain non-negative?

Show answer and solution

For $r=2$ :

k_1\approx-0.693147,\qquad k_2\approx0.240227,\qquad k_3\approx0.306853.

For $r=1/2$ :

k_1\approx0.693147,\qquad k_2\approx0.240227,\qquad k_3\approx0.193147.

Only k1 changes sign. k2 and k3 remain non-negative.

A top-p nucleus

The base distribution is already sorted:

p=(0.50,0.30,0.15,0.05).

Top-p uses threshold $0.80$ , and the token that first reaches the threshold is retained.

Which tokens remain?
What is the renormalized distribution?
What changes at threshold $0.81$ ?

Show answer and solution

At threshold $0.80$ , the first two tokens form the smallest prefix because $0.50+0.30=0.80$ . Renormalization gives

(0.625,0.375,0,0).

At $0.81$ , the third token is needed. Retained mass is $0.95$ and the distribution is

(0.50/0.95,0.30/0.95,0.15/0.95,0).

Truncation correction

Top-p retains a set with base mass $Z_S=0.85$ . The selected token had base probability $0.25$ .

Find its post-renormalization probability and log-probability.
By how much is the new log-probability larger?
By what factor is a four-position sequence ratio wrong if base rather than behavior log-probabilities are used at every position?

Show answer and solution

\widetilde p=0.25/0.85\approx0.294118,

and

\log\widetilde p=\log0.25-\log0.85\approx-1.223775.

The increase is

-\log0.85\approx0.162519.

Across four positions, the multiplicative error is

\exp(4[-\log0.85])=1/0.85^4\approx1.915686.

Whether that factor appears directly or inversely depends on which side of the ratio uses the wrong value.

Support violation

Behavior and target policies over three actions are

\mu=(0.7,0.3,0), \qquad \pi=(0.6,0.2,0.2).

Can $\mathbb{E}_{a\sim\pi}[f(a)]$ be estimated by ordinary ratios $\pi(a)/\mu(a)$ from data generated by $\mu$ ? Which action causes the problem? What can be changed in data collection or in the target policy?

Show answer and solution

No. The third action requires

\pi(3)/\mu(3)=0.2/0,

which is undefined, and behavior data never contains that action. Data collection can use a full-support behavior policy or more exploration, or the target policy can be restricted to behavior support. Clipping a numerical ratio does not create missing information.

Stable `log_softmax`

Implement:

Python

import numpy as np


def stable_log_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    raise NotImplementedError

It must remain stable for logits of order $10^3$ , preserve shape, and normalize probabilities along the selected axis to at least $10^{-12}$ accuracy in float64. Reject NaN and infinite inputs.

Open checks:

Python

logits = np.array([[1000.0, 1001.0, 999.0], [-1000.0, -999.0, -1001.0]])
result = stable_log_softmax(logits, axis=-1)
assert result.shape == logits.shape
np.testing.assert_allclose(np.exp(result).sum(axis=-1), np.ones(2), atol=1e-12, rtol=0.0)
np.testing.assert_allclose(result[0] - result[1], np.zeros(3), atol=1e-12)

Show answer and solution

Python

import numpy as np


def stable_log_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if z.ndim == 0:
        raise ValueError("logits must have at least one dimension")
    if not np.all(np.isfinite(z)):
        raise ValueError("logits must be finite")
    if isinstance(axis, bool) or not isinstance(axis, (int, np.integer)):
        raise TypeError("axis must be an integer")
    axis = int(axis)
    if axis < 0:
        axis += z.ndim
    if not 0 <= axis < z.ndim:
        raise ValueError("axis outside array dimensions")
    maximum = np.max(z, axis=axis, keepdims=True)
    shifted = z - maximum
    log_normalizer = np.log(np.exp(shifted).sum(axis=axis, keepdims=True))
    return shifted - log_normalizer

Subtracting the maximum leaves softmax unchanged while preventing overflow. The log normalizer is then evaluated on non-positive shifted logits.

Causal shift and response mask

Implement a self-contained function that accepts model_logits with shape [B,L,V], input_ids with shape [B,L], and a response-token mask of the same shape. Return:

selected log-probabilities for targets input_ids[:, 1:];
the shifted mask response_token_mask[:, 1:].

Signature:

Python

import numpy as np


def causal_response_logprobs(
    model_logits: np.ndarray,
    input_ids: np.ndarray,
    response_token_mask: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    raise NotImplementedError

Open checks:

Python

logits = np.zeros((1, 3, 4))
input_ids = np.array([[0, 1, 2]])
response_mask = np.array([[0, 1, 1]])
selected, shifted = causal_response_logprobs(logits, input_ids, response_mask)
assert selected.shape == (1, 2) and shifted.shape == (1, 2)
assert np.allclose(selected, -np.log(4.0))
assert np.array_equal(shifted, np.array([[1, 1]], dtype=bool))

Show answer and solution

Python

import numpy as np


def _stable_log_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if not np.all(np.isfinite(z)):
        raise ValueError("logits must be finite")
    maximum = np.max(z, axis=axis, keepdims=True)
    shifted = z - maximum
    return shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))


def causal_response_logprobs(
    model_logits: np.ndarray,
    input_ids: np.ndarray,
    response_token_mask: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    logits = np.asarray(model_logits, dtype=np.float64)
    ids = np.asarray(input_ids, dtype=np.int64)
    mask = np.asarray(response_token_mask, dtype=bool)
    if logits.ndim != 3:
        raise ValueError("model_logits must have shape [B, L, V]")
    if ids.ndim != 2 or ids.shape != logits.shape[:2] or mask.shape != ids.shape:
        raise ValueError("input_ids and response_token_mask must match logits[:2]")
    if logits.shape[1] < 2:
        raise ValueError("the sequence must contain at least two positions")
    vocabulary_size = logits.shape[-1]
    if np.any(ids < 0) or np.any(ids >= vocabulary_size):
        raise ValueError("token id outside vocabulary")

    shifted_logprobs = _stable_log_softmax(logits[:, :-1, :], axis=-1)
    targets = ids[:, 1:]
    selected = np.take_along_axis(shifted_logprobs, targets[..., None], axis=-1)[..., 0]
    return selected, mask[:, 1:]

Logits at position $t$ predict the token at position $t+1$ , so token identifiers and the response mask must be shifted consistently.

Sequence metrics

Given selected log-probabilities and a mask, compute for every response:

the active-token count;
the sum of log-probabilities;
the mean negative log-probability;
perplexity.

An empty mask must raise ValueError. Return an explicitly typed object:

Python

from dataclasses import dataclass
import numpy as np


@dataclass(frozen=True)
class SequenceMetrics:
    token_count: np.ndarray
    sum_logprob: np.ndarray
    mean_nll: np.ndarray
    perplexity: np.ndarray


def masked_sequence_metrics(
    selected_logprobs: np.ndarray,
    mask: np.ndarray,
) -> SequenceMetrics:
    raise NotImplementedError

Open checks:

Python

metrics = masked_sequence_metrics(np.array([[-0.1, -0.2, -1.0]]), np.array([[0, 1, 1]]))
assert metrics.token_count[0] == 2
assert abs(metrics.sum_logprob[0] - (-1.2)) < 1e-12
assert abs(metrics.mean_nll[0] - 0.6) < 1e-12
assert abs(metrics.perplexity[0] - np.exp(0.6)) < 1e-9

Show answer and solution

Python

from dataclasses import dataclass
import numpy as np


@dataclass(frozen=True)
class SequenceMetrics:
    token_count: np.ndarray
    sum_logprob: np.ndarray
    mean_nll: np.ndarray
    perplexity: np.ndarray


def masked_sequence_metrics(
    selected_logprobs: np.ndarray,
    mask: np.ndarray,
) -> SequenceMetrics:
    logprobs = np.asarray(selected_logprobs, dtype=np.float64)
    active = np.asarray(mask, dtype=bool)
    if logprobs.shape != active.shape:
        raise ValueError("selected_logprobs and mask must have the same shape")
    if logprobs.ndim == 1:
        logprobs = logprobs[None, :]
        active = active[None, :]
    elif logprobs.ndim != 2:
        raise ValueError("expected shape [B, T] or [T]")
    if np.any(~np.isfinite(logprobs[active])):
        raise ValueError("active log-probabilities must be finite")

    token_count = active.sum(axis=1).astype(np.int64)
    if np.any(token_count == 0):
        raise ValueError("empty sequence mask")
    sum_logprob = np.where(active, logprobs, 0.0).sum(axis=1)
    mean_nll = -sum_logprob / token_count
    return SequenceMetrics(
        token_count=token_count,
        sum_logprob=sum_logprob,
        mean_nll=mean_nll,
        perplexity=np.exp(mean_nll),
    )

Exact KL and k1–k3 estimators

Implement two self-contained functions:

Python

import numpy as np


def exact_kl_from_logits(p_logits: np.ndarray, q_logits: np.ndarray) -> np.ndarray:
    raise NotImplementedError


def kl_sample_values(log_q: np.ndarray, log_p: np.ndarray) -> dict[str, np.ndarray]:
    raise NotImplementedError

The first returns $D_{\mathrm{KL}}(p\|q)$ along the final axis. In the second, let $r=p/q$ and return sample values k1, k2, and k3.

Open checks:

Python

z = np.array([[1.0, -1.0], [0.5, 0.5]])
np.testing.assert_allclose(exact_kl_from_logits(z, z), np.zeros(2), atol=1e-14)
values = kl_sample_values(np.array([-0.4, -1.2]), np.array([-0.5, -0.7]))
assert set(values) == {"k1", "k2", "k3"}
assert np.all(values["k3"] >= -1e-14)

Show answer and solution

Python

import numpy as np


def _stable_log_softmax(logits: np.ndarray, axis: int = -1) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if not np.all(np.isfinite(z)):
        raise ValueError("logits must be finite")
    maximum = np.max(z, axis=axis, keepdims=True)
    shifted = z - maximum
    return shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))


def exact_kl_from_logits(p_logits: np.ndarray, q_logits: np.ndarray) -> np.ndarray:
    p_logits = np.asarray(p_logits, dtype=np.float64)
    q_logits = np.asarray(q_logits, dtype=np.float64)
    if p_logits.shape != q_logits.shape:
        raise ValueError("p_logits and q_logits must have the same shape")
    log_p = _stable_log_softmax(p_logits, axis=-1)
    log_q = _stable_log_softmax(q_logits, axis=-1)
    p = np.exp(log_p)
    return np.sum(p * (log_p - log_q), axis=-1)


def kl_sample_values(log_q: np.ndarray, log_p: np.ndarray) -> dict[str, np.ndarray]:
    log_q = np.asarray(log_q, dtype=np.float64)
    log_p = np.asarray(log_p, dtype=np.float64)
    if log_q.shape != log_p.shape or not (np.all(np.isfinite(log_q)) and np.all(np.isfinite(log_p))):
        raise ValueError("log_q and log_p must be finite arrays of the same shape")
    log_ratio = log_p - log_q
    ratio = np.exp(log_ratio)
    return {
        "k1": -log_ratio,
        "k2": 0.5 * log_ratio**2,
        "k3": ratio - 1.0 - log_ratio,
    }

Non-negativity of k3 follows from $u-1-\log u\ge 0$ for $u>0$ . A tiny negative value at machine precision may be clipped for display, but a material numerical error should not be hidden.

Temperature, top-k, and top-p

Implement three standalone distribution transforms. Temperature must be positive; top-k and top-p must renormalize the retained support. Use a stable sort so the smaller token identifier wins ties.

Python

import numpy as np


def temperature_distribution(logits: np.ndarray, temperature: float) -> np.ndarray:
    raise NotImplementedError


def top_k_distribution(logits: np.ndarray, k: int) -> np.ndarray:
    raise NotImplementedError


def top_p_distribution(
    logits: np.ndarray,
    top_p: float,
    *,
    min_tokens_to_keep: int = 1,
) -> np.ndarray:
    raise NotImplementedError

Open checks:

Python

logits = np.array([2.0, 1.0, 0.0, -1.0])
warm = temperature_distribution(logits, 2.0)
assert abs(warm[0] - 0.455054) < 1e-6
assert abs(float(warm.sum()) - 1.0) < 1e-12
kept = top_k_distribution(logits, 2)
assert np.allclose(kept[2:], 0.0)
assert abs(kept[0] - 0.731059) < 1e-6
nucleus = top_p_distribution(logits, 0.8)
assert np.allclose(nucleus, kept)

Show answer and solution

Python

import numpy as np


def _softmax_1d(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
    z = np.asarray(logits, dtype=np.float64)
    if z.ndim != 1 or z.size == 0 or not np.all(np.isfinite(z)):
        raise ValueError("logits must be a non-empty finite one-dimensional array")
    if not np.isfinite(temperature) or temperature <= 0.0:
        raise ValueError("temperature must be a positive finite number")
    scaled = z / temperature
    shifted = scaled - np.max(scaled)
    probabilities = np.exp(shifted)
    return probabilities / probabilities.sum()


def _renormalize_support(probabilities: np.ndarray, keep: np.ndarray) -> np.ndarray:
    out = np.zeros_like(probabilities, dtype=np.float64)
    out[keep] = probabilities[keep]
    mass = float(out.sum())
    if mass <= 0.0:
        raise ValueError("the retained support has zero probability mass")
    return out / mass


def temperature_distribution(logits: np.ndarray, temperature: float) -> np.ndarray:
    return _softmax_1d(logits, temperature)


def top_k_distribution(logits: np.ndarray, k: int) -> np.ndarray:
    probabilities = _softmax_1d(logits)
    if isinstance(k, bool) or not isinstance(k, (int, np.integer)) or not (1 <= int(k) <= probabilities.size):
        raise ValueError("k must be an integer between 1 and the vocabulary size")
    order = np.argsort(-probabilities, kind="stable")
    return _renormalize_support(probabilities, order[: int(k)])


def top_p_distribution(
    logits: np.ndarray,
    top_p: float,
    *,
    min_tokens_to_keep: int = 1,
) -> np.ndarray:
    probabilities = _softmax_1d(logits)
    if not np.isfinite(top_p) or not (0.0 < top_p <= 1.0):
        raise ValueError("top_p must lie in (0, 1]")
    if (
        isinstance(min_tokens_to_keep, bool)
        or not isinstance(min_tokens_to_keep, (int, np.integer))
        or not (1 <= int(min_tokens_to_keep) <= probabilities.size)
    ):
        raise ValueError("min_tokens_to_keep must be a valid positive integer")

    order = np.argsort(-probabilities, kind="stable")
    cumulative = np.cumsum(probabilities[order])
    threshold_count = int(np.searchsorted(cumulative, top_p, side="left")) + 1
    keep_count = max(int(min_tokens_to_keep), threshold_count)
    return _renormalize_support(probabilities, order[:keep_count])

When log-probabilities are needed, map zero probabilities in the truncated distribution explicitly to $-\infty$ . These functions test the individual transforms; a complete decoder may compose them in a documented order.

Validate a rollout batch

Implement a validator for

text

response_ids, response_mask,
old_logprobs, ref_logprobs,
sequence_rewards,
terminated, truncated, stop_reason

Check:

aligned token-level shapes;
[B] sequence-level shapes;
a non-empty contiguous response mask;
finite active-token log-probabilities;
NaN on padding positions;
finite rewards;
a non-empty stop_reason string for every response.

Do not impose a universal ban on terminated=True and truncated=True occurring together: a task contract can allow a genuine terminal event to coincide with an external boundary. Structural validation and semantic mapping are separate operations.

Signature:

Python

def validate_rollout_batch(batch: Mapping[str, Any]) -> RolloutBatchSummary:
    ...

Open check:

Python

batch = {
    "response_ids": np.array([[5, 6], [7, 8]]),
    "response_mask": np.array([[1, 1], [1, 0]]),
    "old_logprobs": np.array([[-0.1, -0.2], [-0.3, np.nan]]),
    "ref_logprobs": np.array([[-0.2, -0.3], [-0.4, np.nan]]),
    "sequence_rewards": np.array([1.0, 0.0]),
    "terminated": np.array([True, False]),
    "truncated": np.array([False, True]),
    "stop_reason": np.array(["eos", "length"], dtype=object),
}
summary = validate_rollout_batch(batch)
assert summary.batch_size == 2 and summary.active_tokens == 3
assert summary.terminated_sequences == 1 and summary.truncated_sequences == 1

Show answer and solution

Python

from dataclasses import dataclass
from typing import Any, Mapping
import numpy as np


@dataclass(frozen=True)
class RolloutBatchSummary:
    batch_size: int
    max_response_tokens: int
    active_tokens: int
    terminated_sequences: int
    truncated_sequences: int
    mean_response_length: float


def validate_rollout_batch(batch: Mapping[str, Any]) -> RolloutBatchSummary:
    required = {
        "response_ids", "response_mask", "old_logprobs", "ref_logprobs",
        "sequence_rewards", "terminated", "truncated", "stop_reason",
    }
    missing = sorted(required - set(batch))
    if missing:
        raise ValueError("missing fields: " + ", ".join(missing))

    ids = np.asarray(batch["response_ids"])
    mask = np.asarray(batch["response_mask"], dtype=bool)
    old = np.asarray(batch["old_logprobs"], dtype=np.float64)
    ref = np.asarray(batch["ref_logprobs"], dtype=np.float64)
    rewards = np.asarray(batch["sequence_rewards"], dtype=np.float64)
    terminated = np.asarray(batch["terminated"], dtype=bool)
    truncated = np.asarray(batch["truncated"], dtype=bool)
    reasons = np.asarray(batch["stop_reason"], dtype=object)

    if ids.ndim != 2 or not (ids.shape == mask.shape == old.shape == ref.shape):
        raise ValueError("token-level arrays must align")
    batch_size, max_tokens = ids.shape
    if not (rewards.shape == terminated.shape == truncated.shape == reasons.shape == (batch_size,)):
        raise ValueError("sequence-level arrays must have shape (batch,)")

    lengths = mask.sum(axis=1).astype(int)
    if np.any(lengths <= 0):
        raise ValueError("every response must contain an active token")
    expected = np.arange(max_tokens)[None, :] < lengths[:, None]
    if not np.array_equal(mask, expected):
        raise ValueError("response mask must be a contiguous right-padded prefix")
    if np.any(~np.isfinite(old[mask])) or np.any(~np.isfinite(ref[mask])):
        raise ValueError("active log-probabilities must be finite")
    if np.any(np.isfinite(old[~mask])) or np.any(np.isfinite(ref[~mask])):
        raise ValueError("padding log-probabilities must be NaN")
    if np.any(~np.isfinite(rewards)):
        raise ValueError("rewards must be finite")
    if any(not isinstance(x, str) or not x.strip() for x in reasons.tolist()):
        raise ValueError("every stop_reason must be a non-empty string")

    return RolloutBatchSummary(
        batch_size=int(batch_size),
        max_response_tokens=int(max_tokens),
        active_tokens=int(mask.sum()),
        terminated_sequences=int(terminated.sum()),
        truncated_sequences=int(truncated.sum()),
        mean_response_length=float(lengths.mean()),
    )

The missing ban on coincident flags is intentional. Whether a particular combination is valid under the environment contract belongs in a separate semantic validator that uses both stop_reason and task rules.

Stop reason and task semantics

Is the prefix a sufficient state?

The one-token error

What does PPL say?

“Negative KL”

Support conditions for k3

Which policy generated the token?

A minimal rollout contract

Sequence log-probability and PPL

The KL chain rule on a depth-two tree

k1, k2, and k3

A top-p nucleus

Truncation correction

Support violation

Stable log_softmax

Causal shift and response mask

Sequence metrics

Exact KL and k1–k3 estimators

Temperature, top-k, and top-p

Validate a rollout batch

Stable `log_softmax`