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

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

# %% ../../../nbs/agents/networks/01_critics.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

# Other libraries
from abc import ABC, abstractmethod

# %% ../../../nbs/agents/networks/01_critics.ipynb 6
class CriticNetworkStateAction(BaseNetwork):
    """
    Simple 3-Layer network for the critic function approximator where both state and action are inputs to the network.
    The forward function takes the state and action as input, concatenates them and passes them through the network.
    
    """
    # 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, action):
        state_action = torch.cat((state.float(), action.float()), dim=1)
        features1 = F.relu(self._h1(state_action))
        features2 = F.relu(self._h2(features1))
        q = self._h3(features2)

        return torch.squeeze(q)

