UCBSocketTester.py

class UCBSocketTester( SocketTester ):

    def __init__(self, socket_order, confidence_level = 2.0 ):                  
        super().__init__(socket_order) # create a standard socket tester                
        self.confidence_level = confidence_level # save the confidence_level
        
    def ucb(self, Q, t, n):
        if n == 0: return float('inf')                
        return Q + self.confidence_level * (np.sqrt(np.log(t) / n))         
        
    def select_socket( self, t ):
        """ UCB Socket Selection"""             
        # choose the socket with the current highest mean reward
        return random_argmax([self.ucb(socket.Q, t, socket.n) for socket in self.sockets])