import torch.nn as nn
import torch.nn.functional as F

'''
critic network to learn the values for the state actions pairs
output the joint q values corresponding to all executable actions 
of the current agent, which is a n_actions-dimension vector
'''

class ComaCritic(nn.Module):
	def __init__(self, input_shape, args):
		super(ComaCritic, self).__init__()
		self.args = args
		self.fc1 = nn.Linear(input_shape, args.critic_dim)
		self.fc2 = nn.Linear(args.critic_dim, args.critic_dim)
		self.fc3 = nn.Linear(args.critic_dim, self.args.n_actions)

	def forward(self, inputs):
		x = F.relu(self.fc1(inputs))
		x = F.relu(self.fc2(x))
		q = self.fc3(x)
		return q