def kl(self, prob0, prob1):
        '''
            Method computes KL Divergence of between two probability distributions
            Note: this is D_KL(prob0 || prob1), not D_KL(prob1 || prob0)
        '''
        if len(prob0.size()) == 3:
            prob0 = prob0.view(-1, 2 * self.d)
            prob1 = prob1.view(-1, 2 * self.d)

        mean0 = prob0[:, :self.d]
        std0 = prob0[:, self.d:]
        mean1 = prob1[:, :self.d]
        std1 = prob1[:, self.d:]
        return ((std1 / std0).log()).sum(dim=1) + (
            (std0.pow(2) + (mean0 - mean1).pow(2)) / (2.0 * std1.pow(2))).sum(dim=1) - 0.5 * self.d

            Method computes KL Divergence of between two probability distributions
            Note: this is D_KL(prob0 || prob1), not D_KL(prob1 || prob0)
        
surreal/model/ppo_net.py