UCB1 Average Reward

Problem

Compute an arm's average reward and its UCB1 index.

With n = len(rewards), avg = mean(rewards), and bonus = c * sqrt(log(total + 1) / n), return [avg, avg + bonus]. An empty reward list or a non-positive total returns [0.0, 0.0].

Hint

The bonus term grows with log(total) and shrinks with the arm's pull count.

Test cases

case 1

in: [[1,0.5,0.5],10,1]

expected: [0.6666666666666666,1.5607015788854923]

case 2

in: [[2],100,2]

expected: [2,6.296566311296154]

case 3

in: [[],5,1]

expected: [0,0]

Starter codePython

def ucb1_average(rewards, total, c):
    # Returns [avg, ucb]
    # Your code here
    pass

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