Beta dist policy

beta_policy.py

class BetaHead(nn.Module):
    def __init__(self, in_features, action_size):
        super(BetaHead, self).__init__()

        self.fcc_c0 = nn.Linear(in_features, action_size)
        nn.init.orthogonal_(self.fcc_c0.weight, gain=0.01)
        nn.init.zeros_(self.fcc_c0.bias)

        self.fcc_c1 = nn.Linear(in_features, action_size)
        nn.init.orthogonal_(self.fcc_c1.weight, gain=0.01)
        nn.init.zeros_(self.fcc_c1.bias)

    def forward(self, x):
        c0 = torch.nn.functional.softplus(self.fcc_c0(x)) + 1.
        c1 = torch.nn.functional.softplus(self.fcc_c1(x)) + 1.

        return torch.distributions.Independent(
            torch.distributions.Beta(c1, c0), 1
        )