'Calculate KL between one policy output and another.'

def get_kl(self, my_logits, other_logits):

    kl = []
    for (i, (act_dim, act_type)) in enumerate(self.env_spec.act_dims_and_types):
        sampling_dim = self.env_spec.sampling_dim(act_dim, act_type)
        single_my_logits = my_logits[i]
        single_other_logits = other_logits[i]
        if self.env_spec.is_discrete(act_type):
            my_probs = tf.nn.softmax(single_my_logits)
            my_log_probs = tf.nn.log_softmax(single_my_logits)
            other_log_probs = tf.nn.log_softmax(single_other_logits)
            my_kl = tf.reduce_sum((my_probs * (my_log_probs - other_log_probs)), (-1))
        elif self.env_spec.is_box(act_type):
            my_means = single_my_logits[:, :(sampling_dim / 2)]
            my_std = single_my_logits[:, (sampling_dim / 2):]
            other_means = single_other_logits[:, :(sampling_dim / 2)]
            other_std = single_other_logits[:, (sampling_dim / 2):]
            my_kl = tf.reduce_sum(((tf.log((other_std / my_std)) + ((tf.square(my_std) + tf.square((my_means - other_means))) / (2.0 * tf.square(other_std)))) - 0.5), (-1))
        else:
            assert False
        kl.append(my_kl)
    return kl
