Module 4. Policy gradients: from REINFORCE to PPO

Statements 2 and 4 are correct.

• 2. The identity moves the derivative onto $\log p_\theta(X)$, so $f(X)$ only needs to be evaluated on the sampled outcome.
• 4. A program or human can supply reward. The required differentiable object is the log-probability of the sampled outcome.

Statement 1 is false: sampling remains discrete; the identity constructs a different gradient estimator. Statement 3 is false: Monte Carlo noise remains and may be large. Statement 5 is false: differentiation must be interchangeable with the sum or integral.

With direct dependence, the product rule gives

The score term accounts for the changing distribution; the direct term accounts for the integrand changing at a fixed outcome. Classical REINFORCE assumes environment reward has no direct policy-parameter dependence. A differentiable KL penalty or another parametric regularizer must either retain the direct term or deliberately stop that gradient and thereby define a different computational objective.

1. Yes. A fixed simulator appears in $p_\theta(\tau)$, but its log-probability has zero derivative with respect to $\theta$.

2. No. If transitions share parameter $\theta$, an extra term $\nabla_\theta\log P_\theta(s'\mid s,a)$ appears.

3. Yes, from the explicit derivative. A frozen reward model need not be differentiated, but its output multiplies the score and determines the update direction.

4. No in that estimator. If we deliberately backpropagate through a differentiable world model, transition derivatives are part of the pathwise gradient.

The environment remains in the trajectory distribution. It controls which states and rewards are observed, so “the transition derivative is zero” does not mean “transitions do not affect the gradient estimate through the data.”

Conditioned on the history before $a_t$, past rewards are fixed. Their contribution has the form

Only rewards that the action could influence remain.

If the objective is literally

regrouping the reward terms gives the score weight $\gamma^tG_t$. A formula containing only $G_t$ commonly uses a discounted occupancy measure or another normalization of the objective. It is not the identical algebraic line under an unchanged convention.

Directly valid: 1, 2, and 3.

• A constant and a state-only function do not depend on the current action.
• The mean of other independent responses is conditionally independent of the current response.

Invalid without correction: 4. $Q_\psi(s,a)$ depends on the current action, so the zero-mean score proof does not apply.

Known finite-sample scaling: 5. The same-group mean contains the current reward. Under conditional independence it scales the expected gradient by $(G-1)/G$ rather than preserving it exactly.

1. $\bar R$ contains $R_i/G$, and $R_i$ depends on the current response $Y_i$.

2. The expected gradient scale is $(G-1)/G$.

3. The mean of the other responses is conditionally independent of $Y_i$, so its product with the score has zero expectation.

4. Algebraically,

5. Both advantage constructions return zero. There is no within-group ranking signal.

Case I — external truncation followed by reset.

1. Yes. The mathematical process can continue, so bootstrap from the final observation before reset.

2. No. The next array record after reset belongs to another episode.

3. terminated=False, trace_cut=True.

Case II — a local rollout-buffer boundary.

1. Yes. The same episode continues, so the next value is needed for bootstrap.

2. Not inside this array when the following fragment is not joined to the GAE computation. A streaming implementation could carry recursion state across fragments.

3. terminated=False, trace_cut=True on the final local record.

A single done flag conflates whether continuation value exists with whether the next advantage belongs to the current backward recursion. Zeroing bootstrap at every done biases targets near time limits and buffer boundaries: removing positive continuation lowers the target, while removing negative continuation raises it.

With $\varepsilon=0.2$, the interval is $[0.8,1.2]$.

All four ratios are outside the interval, so ratio_outside_fraction=1. The clipped branch becomes the pessimistic minimum only in rows one and four, hence active_clip_fraction=2/4=0.5.

1. Both use several responses to one prompt and form a relative within-group signal without requiring a separate critic.

2. Basic RLOO uses a leave-one-out mean and a REINFORCE objective. Canonical GRPO uses PPO-like probability ratios and clipping, usually normalizes the group signal, and adds regularization toward a reference policy.

3. The response log-probability is a sum of token log-probabilities, so a whole-response coefficient multiplying that sum is the score function of the full sequence.

4. The coefficient is identical for every token and does not identify which local choice caused the reward. It optimizes the whole-response objective but does not localize causal credit.

First,

For action 2,

The components sum to zero up to rounding.

Backward recursion gives

Thus

Applying the outer discount:

The second vector is the direct score weight for the literal objective whose rewards are discounted from the start of the episode.

The action probabilities are

The state value is

Score vectors:

Squared norms:

Substitution yields

The two baselines differ because $b^*$ weights action values by squared score norm. The rarer second action has a much larger score vector and therefore receives much more weight than it does in $V^\pi$.

The group mean is

Centered advantages:

Leave-one-out baselines:

RLOO advantages:

They are $4/3$ times the centered advantages. The expected scale of the same-group estimator is

The first transition is truncated but not terminal:

The second is terminal:

Because episode_end[0]=1, the recursion does not carry $\delta_1$ backward:

Reward 100 belongs to another episode after reset(). It is adjacent only in storage, not in the trajectory.

The unclipped terms are

The clipped ratios are $(1.2,0.8,1.2,0.8)$, producing

The elementwise pessimistic minimum is

The mean objective is

Every ratio lies outside the interval, so ratio_outside_fraction=1. Clipping changes the minimum in two of four cases, so active_clip_fraction=0.5.

Each $r=1.5$ contributes $0.5-\log1.5\approx0.094535$, while each $r=0.5$ contributes $-0.5-\log0.5\approx0.193147$. Hence

This is a nonnegative sample diagnostic. Its old-policy expectation equals the reverse KL under the required support conditions, but a finite batch does not make it an exact KL constraint.

import numpy as np

def stable_softmax(logits: np.ndarray) -> np.ndarray:
    z = np.asarray(logits, dtype=float)
    if z.ndim != 1:
        raise ValueError("logits must be one-dimensional")
    shifted = z - np.max(z)
    exp_z = np.exp(shifted)
    return exp_z / np.sum(exp_z)

def grad_log_softmax(logits: np.ndarray, action: int) -> np.ndarray:
    p = stable_softmax(logits)
    if not 0 <= action < len(p):
        raise ValueError("action is out of range")
    result = -p
    result[action] += 1.0
    return result

Subtracting the maximum leaves softmax unchanged while preventing exponential overflow.

import numpy as np

def reward_to_go(rewards: np.ndarray, gamma: float) -> np.ndarray:
    r = np.asarray(rewards, dtype=float)
    if r.ndim != 1:
        raise ValueError("rewards must be one-dimensional")
    if not 0.0 <= gamma <= 1.0:
        raise ValueError("gamma must lie in [0, 1]")
    result = np.empty_like(r)
    running = 0.0
    for t in range(len(r) - 1, -1, -1):
        running = r[t] + gamma * running
        result[t] = running
    return result

def discounted_score_weights(rewards: np.ndarray, gamma: float) -> np.ndarray:
    g = reward_to_go(rewards, gamma)
    return (gamma ** np.arange(len(g), dtype=float)) * g

One backward pass gives $O(T)$ time and constant additional storage beyond the result.

import numpy as np

def leave_one_out_baselines(rewards: np.ndarray) -> np.ndarray:
    r = np.asarray(rewards, dtype=float)
    if r.ndim != 1 or len(r) < 2:
        raise ValueError("at least two rewards are required")
    return (np.sum(r) - r) / (len(r) - 1)

def group_mean_advantages(rewards: np.ndarray) -> np.ndarray:
    r = np.asarray(rewards, dtype=float)
    if r.ndim != 1 or len(r) == 0:
        raise ValueError("rewards must be a non-empty vector")
    return r - np.mean(r)

def group_zscore_advantages(rewards: np.ndarray, eps: float = 1e-8) -> np.ndarray:
    centered = group_mean_advantages(rewards)
    scale = float(np.sqrt(np.mean(centered ** 2)))
    if scale < eps:
        return np.zeros_like(centered)
    return centered / scale

This uses the within-group population variance with denominator $G$, a common convention for normalized advantages.

import numpy as np

def gae_advantages(
    rewards: np.ndarray,
    values: np.ndarray,
    next_values: np.ndarray,
    terminated: np.ndarray,
    trace_cuts: np.ndarray,
    gamma: float,
    gae_lambda: float,
) -> np.ndarray:
    r = np.asarray(rewards, dtype=float)
    v = np.asarray(values, dtype=float)
    nv = np.asarray(next_values, dtype=float)
    term = np.asarray(terminated, dtype=bool)
    cuts = np.asarray(trace_cuts, dtype=bool)
    if not (r.shape == v.shape == nv.shape == term.shape == cuts.shape) or r.ndim != 1:
        raise ValueError("all inputs must be aligned one-dimensional arrays")
    if not (0.0 <= gamma <= 1.0 and 0.0 <= gae_lambda <= 1.0):
        raise ValueError("gamma and gae_lambda must lie in [0, 1]")

    deltas = r + gamma * (~term).astype(float) * nv - v
    result = np.empty_like(deltas)
    running = 0.0
    for t in range(len(r) - 1, -1, -1):
        running = deltas[t] + gamma * gae_lambda * (not cuts[t]) * running
        result[t] = running
    return result

The two masks answer different questions. At a truncation or local buffer boundary, terminated=False preserves next-state bootstrap, while trace_cuts=True stops the current batch recursion.

from dataclasses import dataclass
import numpy as np

@dataclass(frozen=True)
class PPODiagnostics:
    ratio_outside_fraction: float
    active_clip_fraction: float
    approx_kl_nonnegative: float
    mean_log_ratio: float

def ppo_clipped_terms(ratio, advantage, clip_epsilon=0.2):
    r = np.asarray(ratio, dtype=float)
    a = np.asarray(advantage, dtype=float)
    if r.shape != a.shape or np.any(r <= 0):
        raise ValueError("positive aligned ratios are required")
    unclipped = r * a
    clipped = np.clip(r, 1 - clip_epsilon, 1 + clip_epsilon) * a
    pessimistic = np.minimum(unclipped, clipped)
    return unclipped, clipped, pessimistic

def ppo_diagnostics(new_log_prob, old_log_prob, advantage, clip_epsilon=0.2):
    log_ratio = np.asarray(new_log_prob, dtype=float) - np.asarray(old_log_prob, dtype=float)
    ratio = np.exp(log_ratio)
    a = np.asarray(advantage, dtype=float)
    unclipped, _, pessimistic = ppo_clipped_terms(ratio, a, clip_epsilon)
    outside = np.mean((ratio < 1 - clip_epsilon) | (ratio > 1 + clip_epsilon))
    active = np.mean(pessimistic < unclipped - 1e-15)
    approx_kl = np.mean((ratio - 1) - log_ratio)
    return PPODiagnostics(float(outside), float(active), float(approx_kl), float(np.mean(log_ratio)))

active_clip_fraction checks whether the pessimistic minimum changed, not merely whether a ratio lies outside the interval.

import numpy as np

def response_token_mask(input_ids, prompt_lengths, pad_token_id, eos_token_id):
    ids = np.asarray(input_ids)
    lengths = np.asarray(prompt_lengths, dtype=int)
    if ids.ndim != 2 or lengths.shape != (ids.shape[0],):
        raise ValueError("input_ids must be [B, T] and prompt_lengths [B]")
    mask = np.zeros((ids.shape[0], ids.shape[1] - 1), dtype=float)
    for b in range(ids.shape[0]):
        active = True
        for pos in range(1, ids.shape[1]):
            token = int(ids[b, pos])
            if pos < lengths[b] or token == pad_token_id or not active:
                continue
            mask[b, pos - 1] = 1.0
            if token == eos_token_id:
                active = False
    return mask

def sequence_log_probability(token_log_probs, mask=None):
    lp = np.asarray(token_log_probs, dtype=float)
    if mask is None:
        return np.sum(lp, axis=-1)
    m = np.asarray(mask, dtype=float)
    if m.shape != lp.shape:
        raise ValueError("mask must match token_log_probs")
    return np.sum(lp * m, axis=-1)

def broadcast_sequence_advantages(sequence_advantages, token_mask):
    a = np.asarray(sequence_advantages, dtype=float)
    m = np.asarray(token_mask, dtype=float)
    if a.ndim != 1 or m.ndim != 2 or m.shape[0] != a.shape[0]:
        raise ValueError("advantages must be [B] and token_mask [B, T]")
    return a[:, None] * m

The mask is shifted relative to the original tokens because logits at position $k$ score the label at position $k+1$. A generated EOS is a policy action and is included; prompt, PAD, and the post-EOS tail are excluded. Broadcasting one advantage implements a whole-response score-function objective, but it does not become a local explanation of token credit.