Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Policy Gradient Methods › Reward-to-Go

Michael BrenndoerferDecember 26, 20256 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 277Policy Gradient Methods › Reward-to-Go

Reading section 9 of 14 · Reward-to-Go · 6 minSection 9 of 14 · 6 min

Section 9 of 14 · 6 min read

Policy Gradient Methods › Reward-to-GoOne section · Sections in this chapter9 / 14 ↓

Sections in this chapter14

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

Another variance reduction technique exploits a basic causal property of sequential decision-making: actions can only affect future rewards, not past ones. A decision made at time $t$ cannot retroactively change what happened at times $1$ through $t-1$ . Yet in basic REINFORCE, we multiply every token's gradient by the full sequence reward, even though early tokens cannot have influenced any "future" rewards from the perspective of later tokens.

Think of a recipe as an analogy. If a chef makes three decisions (which oil to use, which spice to add, whether to reduce the heat), and the final dish receives a good review, it would be wrong to credit the choice of oil for the flavor benefits that the reduced heat provided. The choice of oil could not have influenced whether the sauce reduced properly, since the oil was chosen before the heat decision. Yet basic REINFORCE does exactly this: it credits the oil choice for the reduced-heat flavor improvement.

Instead of multiplying each token's gradient by the total reward, we can use only the rewards that could have been causally influenced by that action. This is called the reward-to-go formulation. The reward-to-go at timestep $t$ sums all rewards from that step forward:

R_t = \sum_{t'=t}^{T} r_{t'}

where:

$R_t$ : reward-to-go from timestep $t$
$r_{t'}$ : immediate reward received at timestep $t'$
The sum starts at the current timestep $t$ and runs to the end of the sequence $T$

This captures only the rewards that could have been influenced by the action at time $t$ , ignoring any rewards that occurred before. The policy gradient with reward-to-go is:

\nabla_\theta J(\theta) = \mathbb{E}_{y \sim \pi_\theta}\left[\sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(y_t|s_t) \cdot R_t\right]

where:

$\mathbb{E}_{y \sim \pi_\theta}$ : expectation over sequences sampled from the policy
$y_t$ : the token generated at timestep $t$
$R_t$ : the cumulative reward from time $t$ onwards (reward-to-go), ignoring past rewards

This does not change the expected gradient (the proof is analogous to the baseline proof: terms involving rewards from before time $t$ vanish in expectation because they are independent of the action at time $t$ ) but reduces variance by removing terms that add noise without giving useful credit assignment signal.

In[17]:

Code

import torch


def compute_rewards_to_go(rewards, gamma=1.0):
    """
    Compute reward-to-go for each timestep.

    Args:
        rewards: List of rewards at each timestep
        gamma: Discount factor (1.0 = no discounting)

    Returns:
        Tensor of rewards-to-go
    """
    T = len(rewards)
    rewards_to_go = torch.zeros(T)

    # Work backwards from the end of the sequence
    running_sum = 0
    for t in reversed(range(T)):
        running_sum = rewards[t] + gamma * running_sum
        rewards_to_go[t] = running_sum

    return rewards_to_go

The backward pass through the sequence is computationally efficient. We start from the last timestep (where reward-to-go equals just the immediate reward) and accumulate backward. The discount factor $\gamma$ down-weights rewards that are further in the future, introducing a form of temporal discounting. In standard RLHF with a single end-of-sequence reward, $\gamma = 1.0$ is common since all rewards are already at the end.

In[18]:

Code

# Example: sequence with only a final reward
intermediate_rewards = [0, 0, 0, 0, 1.0]  # Only final reward is non-zero
rewards_to_go = compute_rewards_to_go(intermediate_rewards)

Out[19]:

Console

Timestep | Immediate Reward | Reward-to-Go
---------------------------------------------
    0    |       0.0        |     1.0
    1    |       0.0        |     1.0
    2    |       0.0        |     1.0
    3    |       0.0        |     1.0
    4    |       1.0        |     1.0

In this example with only a final reward, reward-to-go equals the final reward at all timesteps. This makes sense: from every position in the sequence, the only future reward available is the one at the end. The real benefit of reward-to-go emerges when we have intermediate rewards, such as the per-token KL divergence penalty that RLHF systems often add.

Out[20]:

Visualization

Bar chart of per-token immediate rewards with a highlighted positive first reward bar and small negative bars elsewhere. — Per-token immediate rewards for a sequence with a large positive reward at the first timestep and small negative KL penalties at subsequent timesteps. The front-loaded positive reward is marked in green.

Bar chart of cumulative reward-to-go at each timestep, with bars decreasing toward the end. — Cumulative reward-to-go $R_t$ at each timestep, showing that early tokens see larger future rewards than late tokens. Using reward-to-go instead of the total sequence reward reduces gradient variance by excluding past rewards that could not have been influenced by each action.

In the typical RLHF setup where reward is only given at the end, reward-to-go equals the final reward at all timesteps. However, this structure becomes important when we add auxiliary rewards or penalties, such as the KL divergence penalty we will introduce in subsequent chapters. That penalty applies at each token, creating a richer reward structure where reward-to-go provides meaningful variance reduction by correctly attributing per-token contributions to the overall gradient.

Combining Baselines and Reward-to-GoLink Copied

The two techniques can be combined. We can use a per-token value function $V(s_t)$ as a baseline for each timestep, subtracting it from the reward-to-go at that position:

A_t = R_t - V(s_t)

where:

$A_t$ : the per-token advantage at timestep $t$
$R_t$ : the reward-to-go from timestep $t$
$V(s_t)$ : the value function predicting expected future reward from state $s_t$

This per-token advantage is what actor-critic algorithms compute. The actor is the policy network being trained; the critic is the value function giving the baseline. By training the critic to accurately predict expected future rewards, we get better baselines, which gives us lower-variance advantage estimates, which gives us more stable policy updates.