RLHF Reward Shaping
Implement reward shaping for RLHF with KL penalty.
The shaped reward at each token position t is: Rt​={−β⋅klt​rreward​−β⋅klt​​if t<Tif t=T​
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: Gt​=∑k=0T−t​γkRt+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:
0.1 0.99 2.0 3 0.5 0.3 0.1
-0.0500 -0.0300 1.9900 1.9008 1.9706 1.9900
- First, we calculate the shaped rewards Rt​ for each position t using the given formula:
- For t<T, Rt​=−β⋅klt​=−0.1⋅klt​
- For t=T, RT​=rreward​−β⋅klT​=2.0−0.1⋅0.1=1.99
- This results in R1​=−0.1⋅0.5=−0.05, R2​=−0.1⋅0.3=−0.03, R3​=2.0−0.1⋅0.1=1.99
- Then, we calculate the discounted returns Gt​ for each position t using the formula: Gt​=∑k=0T−t​γkRt+k​
- For t=1, G1​=R1​+γR2​+γ2R3​=−0.05+0.99⋅(−0.03)+0.992⋅1.99
- For t=2, G2​=R2​+γR3​=−0.03+0.99⋅1.99
- For t=3, G3​=R3​=1.99
- The calculated shaped rewards are −0.05, −0.03, 1.99 and the calculated discounted returns are −0.05+0.99⋅(−0.03)+0.992⋅1.99≈1.9008, −0.03+0.99⋅1.99≈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
More from LLM 2: Training & Alignment
Background Knowledge
The problem involves reward shaping for Reinforcement Learning from Human Feedback (RLHF), which is a technique used to train language models. The goal of RLHF is to align the model's behavior with human preferences. The KL penalty is a regularization term that encourages the model to stay close to a reference distribution. In this case, the reference distribution is represented by π_ref(token_t|context), and the model's distribution is π_θ(token_t|context). The KL divergence measures the difference between these two distributions.
The problem also involves discounted returns, which is a concept from reinforcement learning. The idea is to calculate the expected cumulative reward that an agent can receive from a particular state. The discount factor γ determines how much the agent values immediate rewards versus future rewards. A discount factor close to 1 means the agent values future rewards more, while a discount factor close to 0 means the agent values immediate rewards more.
The shaped reward is a way to modify the original reward signal to encourage the model to behave in a certain way. In this case, the shaped reward is a combination of the terminal reward and the KL penalty. The shaped reward is calculated at each token position t, and then the discounted returns are calculated using the shaped rewards.
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.