epsilon_greedy_with_softmax.py

import numpy as np
import torch
from tools.parameter_scheduler import ParameterScheduler
from agents.policies.base_policy import Policy
import random
from tools.rl_constants import Action


class EpsilonGreedySoftmaxPolicy(Policy):
    def __init__(self, action_size: int, epsilon_scheduler: ParameterScheduler, seed: int = None):
        super().__init__(action_size=action_size)
        self.epsilon_scheduler = epsilon_scheduler

        self.action_size = action_size

        # Initialize epsilon
        self.epsilon = self.epsilon_scheduler.initial

        if seed:
            self.set_seed(seed)

    def step(self, episode_number: int):
        self.epsilon = self.epsilon_scheduler.get_param(episode_number)
        return True

    def get_action(self, state: np.array, model: torch.nn.Module) -> Action:
        """ Implement this function for speed"""

        def _get_action_values():
            model.eval()
            with torch.no_grad():
                action_values = model.forward(state, act=True)
            model.train()
            return action_values

        if self.training:
            action_values_ = _get_action_values()
            if random.random() > self.epsilon:
                action = action_values_.max(1)[1].data[0]
            else:
                probs = torch.nn.functional.softmax(action_values_)
                action = np.random.choice(np.arange(0, self.action_size), p=probs.view(-1).numpy())
        else:
            action_values_ = _get_action_values()
            action = action_values_.max(1)[1].data[0]

        return Action(value=action)

    def get_deterministic_policy(self, state_action_values_dict: dict):
        deterministic_policy = {}
        for state in state_action_values_dict:
            deterministic_policy[state] = np.argmax(state_action_values_dict[state])
        return deterministic_policy