Distributional Return Variance

Problem

Compute the variance of a categorical return distribution.

With mean = sum probs[i] * atoms[i], return sum probs[i] * (atoms[i] - mean)^2. An empty distribution returns 0.0.

Hint

Distributional RL predicts the whole distribution, and its spread is a risk signal.

case 1

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

expected: 0.6100000000000001

case 2

in: [[1],[1]]

expected: 0

case 3

in: [[],[]]

expected: 0

case 4

in: [[0,10],[0.5,0.5]]

expected: 25

Starter codePython

def distributional_variance(atoms, probs):
    # Your code here
    pass

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