KL Coefficient Schedule

Problem

Linearly schedule the KL coefficient over training: beta_start + (beta_end - beta_start) * min(1, step / total_steps), with progress clamped to [0, 1]. A non-positive total_steps returns beta_end.

Hint

Warming the KL coefficient up often stabilizes the early phase of RLHF.

Test cases

case 1

in: [0,100,0,0.1]

expected: 0

case 2

in: [50,100,0,0.1]

expected: 0.05

case 3

in: [200,100,0,0.1]

expected: 0.1

case 4

in: [10,0,0.5,0.2]

expected: 0.2

Starter codePython

def kl_coefficient_schedule(step, total_steps, beta_start, beta_end):
    # Your code here
    pass

4 test cases · run your code against them in the editor