| import numpy as np | |
| import torch | |
| class Agent: | |
| def __init__(self, | |
| net, | |
| action_space=None, | |
| exploration_initial_eps=None, | |
| exploration_decay=None, | |
| exploration_final_eps=None): | |
| self.net = net | |
| self.action_space = action_space | |
| self.exploration_initial_eps = exploration_initial_eps | |
| self.exploration_decay = exploration_decay | |
| self.exploration_final_eps = exploration_final_eps | |
| self.epsilon = 0. | |
| def __call__(self, state, device=torch.device('cpu')): | |
| if np.random.random() < self.epsilon: | |
| action = self.get_random_action() | |
| else: | |
| action = self.get_action(state, device) | |
| return action | |
| def get_random_action(self): | |
| action = self.action_space.sample() | |
| return action | |
| def get_action(self, state, device=torch.device('cpu')): | |
| if not isinstance(state, torch.Tensor): | |
| state = torch.tensor([state]) | |
| if device.type != 'cpu': | |
| state = state.cuda(device) | |
| q_values = self.net.eval()(state) | |
| _, action = torch.max(q_values, dim=1) | |
| return int(action.item()) | |
| def update_epsilon(self, step): | |
| self.epsilon = max( | |
| self.exploration_final_eps, self.exploration_final_eps + | |
| (self.exploration_initial_eps - self.exploration_final_eps) * | |
| self.exploration_decay**step) | |
| return self.epsilon | |