agent.py

class Agent:
    
    #HELPFUL FUNCTIONS
    def possible_actions(self, velocity):
        '''
        *** Performs two tasks, can be split up ***
        Universe of actions:  α = [(-1,-1),(-1,0),(0,-1),(-1,1),(0,0),(1,-1),(0,1),(1,0),(1,1)]
                            
        Uses constraints to filter out invalid actions given the velocity
        
        0 <= v_x < 5
        0 <= v_y < 5
        v_x and v_y cannot be made both zero (you can't take an action which would make them zero simultaneously)
        Returns list of possible actions given the velocity
        '''
        α = [(-1,-1),(-1,0),(0,-1),(-1,1),(0,0),(1,-1),(0,1),(1,0),(1,1)]
        α = [np.array(x) for x in α]

        β = []
        for i,x in zip(range(9),α):
            new_vel = np.add(velocity,x)
            if (new_vel[0] < 5) and (new_vel[0] >= 0) and (new_vel[1] < 5) and (new_vel[1] >= 0) and ~(new_vel[0] == 0 and new_vel[1] == 0):
                β.append(i)
        β = np.array(β)
        
        return β
    
    def map_to_1D(self,action):
        α = [(-1,-1),(-1,0),(0,-1),(-1,1),(0,0),(1,-1),(0,1),(1,0),(1,1)]
        for i,x in zip(range(9),α):
            if action[0]==x[0] and action[1]==x[1]:
                return i
    
    def map_to_2D(self,action):
        α = [(-1,-1),(-1,0),(0,-1),(-1,1),(0,0),(1,-1),(0,1),(1,0),(1,1)]
        return α[action]
    
    #CONSTRUCTOR
    def __init__(self):
        pass
    
    def get_action(self, state, policy):
        '''
        Returns action given state using policy
        '''
        return self.map_to_2D(policy(state, self.possible_actions(state[2:4])))