Softmax Temperature Policy

Problem

Return the probability that a softmax policy with the given temperature assigns to one action index.

Compute softmax(logits / temperature) with the max-subtraction trick and return the entry at index. Temperature is positive.

Hint

Low temperature sharpens the distribution; high temperature flattens it.

case 1

in: [[1,2,3],2,1]

expected: 0.6652409557748218

case 2

in: [[0,0],0,1]

expected: 0.5

case 3

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

expected: 0.5

case 4

in: [[5,5,5],0,2]

expected: 0.3333333333333333

Starter codePython

import math
def softmax_action_prob(logits, index, temperature):
    # Your code here
    pass

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