Calculate only value head for given state

def value(self, observation):
        """ Calculate only value head for given state """
        policy_params, q = self(observation)

        # Expectation of Q value with respect to action
        return (torch.exp(policy_params) * q).sum(dim=1)