Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

PPO for Language Models › Reward Assignment in Sequential Generation

Michael BrenndoerferDecember 28, 20255 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 279PPO for Language Models › Reward Assignment in Sequential Generation

Reading section 5 of 15 · Reward Assignment in Sequential Generation · 5 minSection 5 of 15 · 5 min

Section 5 of 15 · 5 min read

PPO for Language Models › Reward Assignment in Sequential GenerationOne section · Sections in this chapter5 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Reward models and RL algorithms use rewards differently, and this mismatch creates one of the central challenges for PPO in language models. The reward model, as we discussed in the chapter on Reward Modeling, takes a complete prompt-response pair and outputs a single scalar score. It evaluates the response as a whole, considering whether it is helpful and coherent as well as accurate and safe. But PPO operates on trajectories with per-timestep rewards. The algorithm expects to receive a reward signal at each step, allowing it to compute advantages and update the policy accordingly.

Think of it this way: the reward model is like a book critic who reads an entire novel and then gives it one star rating. PPO is like a writing coach who wants feedback on every sentence so it can tell the author what to keep and what to cut. The critic's overall assessment captures something real, but the writing coach needs more granular guidance. We have to translate the critic's single rating into per-sentence feedback, which inevitably involves some imprecision.

The standard approach assigns the reward to the final token:

r_t = \begin{cases} R(x, y) & \text{if } t = T \text{ (final token)} \\ 0 & \text{otherwise} \end{cases}

where:

$r_t$ : the reward assigned at step $t$
$R(x, y)$ : the scalar score from the reward model for the complete response
$T$ : the length of the generated sequence (final step)
$x$ : the input prompt
$y$ : the complete generated response

This approach reflects the reality that we only know how good a response is once it is complete. We cannot reliably evaluate a partial response because quality often depends on the full content, including how the response concludes. A response that starts brilliantly but ends poorly deserves a low score, but you would not know that until you reach the end. So we wait until generation finishes, compute the reward, and assign it to the final step.

Sparse rewards make credit assignment difficult. The policy must learn which of its many token choices contributed to the final reward. PPO addresses this through its advantage estimation, but the challenge remains significant. A response might receive a low reward because of a single poor word choice, but that signal must propagate back through dozens of preceding tokens. The value function must learn to predict, from any intermediate state, what the expected final reward will be. This prediction is what allows advantages to differentiate between tokens: some tokens lead to states with high expected rewards, others to states with lower expectations.

In[9]:

Code

import numpy as np


def assign_rewards(response_length, final_reward):
    """
    Standard reward assignment: all reward at final step.

    Args:
        response_length: Number of tokens in the response
        final_reward: Scalar reward from the reward model

    Returns:
        Array of per-token rewards
    """
    rewards = np.zeros(response_length)
    rewards[-1] = final_reward  # Assign all reward to final token
    return rewards


def compute_returns(rewards, gamma=1.0):
    """
    Compute discounted returns from rewards.

    With gamma=1.0 (common in RLHF), each token's return
    equals the final reward.
    """
    T = len(rewards)
    returns = np.zeros(T)
    running_return = 0

    for t in reversed(range(T)):
        running_return = rewards[t] + gamma * running_return
        returns[t] = running_return

    return returns

In[10]:

Code

# Demonstrate reward assignment
response_len = 8
reward = 0.75

rewards = assign_rewards(response_len, reward)
returns = compute_returns(rewards)

Out[11]:

Console

Per-token rewards: [0.   0.   0.   0.   0.   0.   0.   0.75]
Per-token returns: [0.75 0.75 0.75 0.75 0.75 0.75 0.75 0.75]

Every token receives return = 0.75 (the final reward)

Out[12]:

Visualization

Bar chart of per-token rewards for an 8-token response. Seven bars are at zero and the final bar reaches 0.75. — Sparse terminal rewards from the reward model: the final token receives the full score while all preceding tokens receive zero reward, which reflects that only a complete response can be meaningfully evaluated.

Bar chart of per-token returns for an 8-token response with gamma equals 1.0. All eight bars are the same height at 0.75. — Constant returns with discount factor gamma equal to 1.0: every token in the sequence shares the same return equal to the terminal reward, propagating the signal uniformly backward through the sequence.

With a discount factor of $\gamma = 1.0$ , which is standard in RLHF, every token in the sequence receives the same return equal to the final reward. This uniform signal might seem uninformative at first glance. If every token gets the same return, how can the algorithm distinguish good tokens from bad ones? The answer lies in the advantage function, which compares the actual return to the expected return under the value function. A token that leads to a higher-than-expected return receives a positive advantage, while a token that leads to a lower-than-expected return receives a negative advantage. This differential signal, created by the value function's predictions, is what enables learning even with sparse rewards.