C51 Projection Step

Problem

Project a categorical value distribution onto a fixed support for C51.

For each atom compute the target value reward + gamma * support[i] (or just reward for a terminal transition), then distribute its probability to the two neighboring atoms by linear interpolation. Mass below v_min goes to the first atom and mass above v_max to the last. Return the projected probability vector.

Hint

Projection keeps the distribution on the fixed support after the Bellman update.

Test cases

case 1

in: [[0,1,2,3],[0.25,0.25,0.25,0.25],1,0.9,false,0,3]

expected: [0,0.275,0.275,0.44999999999999996]

case 2

in: [[0,1,2,3],[0.1,0.2,0.3,0.4],1,0.5,false,0,3]

expected: [0,0.2,0.6000000000000001,0.2]

case 3

in: [[0,1,2,3],[0.25,0.25,0.25,0.25],2,0.9,true,0,3]

expected: [0,0,1,0]

Starter codePython

def c51_projection(support, probs, reward, gamma, done, v_min, v_max):
    # Your code here
    pass

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