'Calculate KL of distribution with itself.
Used layer only for the gradients.'

def self_kl(self, logits, sampling_dim, act_dim, act_type):

    if self.env_spec.is_discrete(act_type):
        probs = tf.nn.softmax(logits)
        log_probs = tf.nn.log_softmax(logits)
        self_kl = tf.reduce_sum((tf.stop_gradient(probs) * (tf.stop_gradient(log_probs) - log_probs)), (-1))
    elif self.env_spec.is_box(act_type):
        means = logits[:, :(sampling_dim / 2)]
        std = logits[:, (sampling_dim / 2):]
        my_means = tf.stop_gradient(means)
        my_std = tf.stop_gradient(std)
        self_kl = tf.reduce_sum(((tf.log((std / my_std)) + ((tf.square(my_std) + tf.square((my_means - means))) / (2.0 * tf.square(std)))) - 0.5), (-1))
    else:
        assert False
    return self_kl
