# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../../nbs/agents/rl_agents/pre_specified_agents/32_SAC_gumbel.ipynb.

# %% auto 0
__all__ = ['SACGumbel']

# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/32_SAC_gumbel.ipynb 4
# General libraries
import numpy as np

# Torch
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# Networks
from ...networks.critics import CriticNetworkStateAction
from ...networks.actors import ActorNetwork

# Algorithms
from ..sac_gumbel import SAC_gumbel

# Processors
from ...processors.processors import  OneHotToDiscrete


# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/32_SAC_gumbel.ipynb 5
class SACGumbel():

    """
    Soft Actor Critic (SAC) agent for discrete action space. Follows the implementation from https://arxiv.org/pdf/1910.07207.pdf

    Args:
        mdp_info (MDPInfo): Contains relevant information about the environment.
        num_actions (int): Number of actions of the environment. If None, it will be obtained from the environment.
        learning_rate_actor (float): Learning rate for the actor.
        learning_rate_critic (float): Learning rate for the critic.
        learning_rate_alpha (float): Learning rate for the temperature parameter.
        initial_replay_size (int): Number of transitions to save in the replay buffer during startup.
        max_replay_size (int): Maximum number of transitions to save in the replay buffer.
        batch_size (int): Number of transitions to sample each time experience is replayed.
        n_features (int): Number of features for the hidden layers of the networks.
        warmup_transitions (int): Number of transitions to replay before starting to update the networks.
        lr_alpha (float): Learning rate for the temperature parameter.
        tau (float): Parameter for the soft update of the target networks.
        target_entropy (float): Target entropy - follows rule of thumb of 0.98 * -log(1/|A|) if None, otherwise it uses the provided value.
        temperature (float): Temperature of the Gumbel Softmax distribution for backpropagation.
        optimizer (torch.optim): Optimizer to use for the networks.
        squeeze_output (bool): Whether to squeeze the output of the actor network or not.
        use_cuda (bool): Whether to use CUDA or not. If True and not available, it will use CPU.
        agent_name (str): Name of the agent. If set to None will use some default name.

    """

    def __init__(
            self,
            mdp_info,
            num_actions = None,
            learning_rate_actor = 3e-4,
            learning_rate_critic = None,
            learning_rate_alpha = 3e-4,
            initial_replay_size = 64,
            max_replay_size = 50000,
            batch_size = 64,
            n_features = 64,
            warmup_transitions = 100,
            lr_alpha = 3e-4,
            tau = 0.005,
            target_entropy = None,
            temperature = 1.0,
            optimizer = optim.Adam,
            squeeze_output = True,
            use_cuda = True,
            agent_name = None):
        
        use_cuda = use_cuda and torch.cuda.is_available()

        actor_input_shape = mdp_info.observation_space.shape
        
        if num_actions == None:
            actor_output_shape = (mdp_info.action_space.shape[0],)
        else:
            actor_output_shape = (num_actions,)

        critic_input_shape = (actor_input_shape[0] + actor_output_shape[0],)

        if learning_rate_critic is None:
            learning_rate_critic = learning_rate_actor

        if target_entropy is None:
            target_entropy = -np.log((1.0 / mdp_info.action_space.shape[0])) * 0.98

        actor_params = dict(network=ActorNetwork,
                                n_features=n_features,
                                input_shape=actor_input_shape,
                                output_shape=actor_output_shape,
                                use_cuda=use_cuda)

        actor_optimizer = {'class': optimizer,
                    'params': {'lr': learning_rate_alpha}} 
        
        critic_params = dict(network=CriticNetworkStateAction,
                                optimizer={'class': optim.Adam,
                                        'params': {'lr': 3e-4}}, # default: 3e-4
                                loss=F.mse_loss,
                                n_features=n_features,
                                input_shape=critic_input_shape,
                                output_shape=(1,),
                                squeeze_output=squeeze_output,
                                use_cuda=use_cuda)
    
        # Agent
        self.agent = SAC_gumbel(mdp_info, actor_params,
                    actor_optimizer, critic_params, batch_size, initial_replay_size,
                    max_replay_size, warmup_transitions, tau, lr_alpha, temperature = temperature,
                    critic_fit_params=None, target_entropy=target_entropy)
        
        if agent_name is None:
            self.agent.name = 'SAC_gumbel'
        else:
            self.agent.name = agent_name

        self.agent.add_postprocessor(OneHotToDiscrete())

    def __getattr__(self, attr):
        return getattr(self.agent, attr)


