KL Trust Region Step

Problem

Perform a KL-constrained update on the probability simplex.

Compute q_i proportional to p_i * exp(-lr * g_i), normalize q to sum to 1, then return [q, kl] where kl = sum_i p_i * log(p_i / q_i) over entries with p_i > 0. This is the closed form of a linear objective with a KL trust region.

Hint

The exponential tilt is the prox of the negative entropy, and the KL measures how far it moved.

Test cases

case 1

in: [[0.5,0.5],[1,-1],0.1]

expected: [[0.4501660026875221,0.549833997312478],0.004991688821646446]

case 2

in: [[0.2,0.3,0.5],[0,0,0],1]

expected: [[0.2,0.3,0.5],0]

case 3

in: [[0.9,0.1],[-2,1],0.5]

expected: [[0.9758075451312032,0.02419245486879684],0.06912938372077355]

Starter codePython

def kl_constrained_step(p, g, lr):
    # Returns [q, kl]
    # Your code here
    pass

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