import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, value_coef):
        super(Model, self).__init__()
        self.value_coef = value_coef

    def forward(self, log_probs: torch.Tensor, values: torch.Tensor, returns: torch.Tensor,
                advantages: torch.Tensor) -> torch.Tensor:
        actor_loss = -(log_probs * advantages).mean()
        critic_loss = ((values - returns) ** 2).mean()
        loss = actor_loss + self.value_coef * critic_loss
        return loss


batch_size = 256


def get_inputs():
    log_probs = torch.randn(batch_size)
    values = torch.randn(batch_size)
    returns = torch.randn(batch_size)
    advantages = torch.randn(batch_size)
    return [log_probs, values, returns, advantages]


def get_init_inputs():
    value_coef = torch.tensor(0.5)
    return [value_coef]