UCB Action Selection

Choose an arm in a multi-armed bandit with the UCB1 rule.

A bandit has several arms, each paying out from an unknown distribution. You must balance trying arms you know little about (exploration) against pulling the arm that looks best so far (exploitation). The upper confidence bound rule adds a bonus to each arm's average reward. The bonus is large for arms that have been pulled rarely, and it shrinks as an arm is pulled more:

score(a) = values[a] + c * sqrt(ln(t) / counts[a])

Here values[a] is the average reward of arm a so far, counts[a] is how many times it was pulled, and t is the total number of pulls so far.

If any arm has never been pulled, return the first such arm. Otherwise return the arm with the highest score, with ties going to the lower index.

Function signature

def ucb_action(counts, values, t, c=2.0):

Example

Input

ucb_action(np.array([10, 5, 1]), np.array([0.6, 0.5, 0.2]), t=16, c=1.0)

Output

2

Reasoning

ln(16) = 2.773. The bonuses are sqrt(2.773 / 10) = 0.527, sqrt(2.773 / 5) = 0.745 and sqrt(2.773 / 1) = 1.665. The scores are 1.127, 1.245 and 1.865. Arm 2 has the worst average, but it was pulled only once, so it is the most worth checking.

Constraints and notes

• counts holds non-negative integers and t >= 1.
• Return a plain int.