RLHF Reward Shaping

hardPROPreference OptimizationLLM 2: Training & Alignment

Implement reward shaping for RLHF with KL penalty.

The shaped reward at each token position t is: $R_t = \begin{cases} -\beta \cdot \text{kl}_t & \text{if } t < T \\ r_{\text{reward}} - \beta \cdot \text{kl}_t & \text{if } t = T \end{cases}$

where kl_t = log(π_θ(token_t|context)) - log(π_ref(token_t|context)) is the per-token KL, r_reward is the terminal reward from the reward model, β is the KL penalty coefficient, and T is the last token position.

Then compute the discounted returns (rewards-to-go) for each position: $G_t = \sum_{k=0}^{T-t} \gamma^k R_{t+k}$

Input:

Line 1: beta gamma (KL coefficient, discount factor)
Line 2: r_reward (terminal reward)
Line 3: T (sequence length)
Line 4: space-separated per-token KL values (T values)

Output:

Line 1: Shaped rewards for each position, rounded to 4 decimal places
Line 2: Discounted returns for each position, rounded to 4 decimal places

Example

Input

0.1 0.99
2.0
3
0.5 0.3 0.1

Output

-0.0500 -0.0300 1.9900
1.9008 1.9706 1.9900

Reasoning

First, we calculate the shaped rewards $R_t$ $R_{t}$ for each position $t$ $t$ using the given formula:
- For $t < T$ , $R_t = -\beta \cdot \text{kl}_t = -0.1 \cdot \text{kl}_t$
- For $t = T$ , $R_T = r_{\text{reward}} - \beta \cdot \text{kl}_T = 2.0 - 0.1 \cdot 0.1 = 1.99$
- This results in $R_1 = -0.1 \cdot 0.5 = -0.05$ , $R_2 = -0.1 \cdot 0.3 = -0.03$ , $R_3 = 2.0 - 0.1 \cdot 0.1 = 1.99$
Then, we calculate the discounted returns $G_t$ $G_{t}$ for each position $t$ $t$ using the formula: $G_t = \sum_{k=0}^{T-t} \gamma^k R_{t+k}$ $G_{t} = \sum_{k = 0}^{T - t} γ^{k} R_{t + k}$
- For $t = 1$ , $G_1 = R_1 + \gamma R_2 + \gamma^2 R_3 = -0.05 + 0.99 \cdot (-0.03) + 0.99^2 \cdot 1.99$
- For $t = 2$ , $G_2 = R_2 + \gamma R_3 = -0.03 + 0.99 \cdot 1.99$
- For $t = 3$ , $G_3 = R_3 = 1.99$
The calculated shaped rewards are $-0.05$ , $-0.03$ , $1.99$ and the calculated discounted returns are $-0.05 + 0.99 \cdot (-0.03) + 0.99^2 \cdot 1.99 \approx 1.9008$ , $-0.03 + 0.99 \cdot 1.99 \approx 1.9706$ , $1.99$
The final output is the shaped rewards and discounted returns rounded to 4 decimal places: $-0.0500$ , $-0.0300$ , $1.9900$ and $1.9008$ ,

Constraints

0 < beta <= 1, 0 < gamma <= 1
1 <= T <= 20
Round to 4 decimal places

Learn this in

Safety & EthicsLLM Study Plan RLHF & AlignmentLLM Study Plan Text Generation & SummarizationNLP Study Plan Training Recipes: Pretraining, SFT & RLHFVLM Study Plan

More from LLM 2: Training & Alignment

Implicit Reward from DPOmedium PPO Clipped Objectivemedium DPO Lossmedium KL Divergence Penaltyeasy Best-of-N Samplinghard

RLHF Reward Shaping

Example

Constraints

Learn this in

More from LLM 2: Training & Alignment

Background Knowledge

Continue the full explanation

Editor locked