'Calculate KL between one policy and another on batch of episodes.'

def calculate_kl(self, my_logits, other_logits):

    batch_size = tf.shape(my_logits[0])[1]
    time_length = tf.shape(my_logits[0])[0]
    reshaped_my_logits = [tf.reshape(my_logit, [(batch_size * time_length), (-1)]) for my_logit in my_logits]
    reshaped_other_logits = [tf.reshape(other_logit, [(batch_size * time_length), (-1)]) for other_logit in other_logits]
    kl = self.get_kl(reshaped_my_logits, reshaped_other_logits)
    kl = [tf.reshape(kkl, [time_length, batch_size]) for kkl in kl]
    return kl
