# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/agents/networks/03_actors.ipynb.

# %% auto 0
__all__ = ['ActorNetwork', 'ActorNetworkDiscrete']

# %% ../../../nbs/agents/networks/03_actors.ipynb 4
# Networks
from .base import BaseNetwork

# General libraries
import numpy as np

# Pytorch
import torch
import torch.nn as nn
import torch.nn.functional as F

# %% ../../../nbs/agents/networks/03_actors.ipynb 6
class ActorNetwork(BaseNetwork):
    """
    Simple 3-Layer network for the actor function approximator.
    
    """
    # TODO make the network architecture more flexible

    def __init__(self, input_shape, output_shape, n_features, **kwargs):
        super().__init__(input_shape, output_shape, n_features, **kwargs)

    def forward(self, state):
        
        #features1 = F.relu(self._h1(torch.squeeze(state, 1).float()))
        features1 = F.relu(self._h1(state.float()))
        features2 = F.relu(self._h2(features1))
        a = self._h3(features2)

        return a

# %% ../../../nbs/agents/networks/03_actors.ipynb 7
class ActorNetworkDiscrete(BaseNetwork):
    """
    Simple 3-Layer network for the actor function approximator. Adding a softmax layer to get a probability distribution over the actions.
    
    """
    # TODO make the network architecture more flexible

    def __init__(self, input_shape, output_shape, n_features, **kwargs):
        super().__init__(input_shape, output_shape, n_features, **kwargs)

    def forward(self, state):
        #features1 = F.relu(self._h1(torch.squeeze(state, 1).float()))
        features1 = F.relu(self._h1(state.float()))
        features2 = F.relu(self._h2(features1))
        a = self._h3(features2)
        a = F.softmax(a, dim=1)

        return a
