def kl_divergence(policy, all_obs, old_dist):
    """
    Compute the kl divergence of the old and new dist.

    Parameters
    ----------
    policy (nn.Module):
    all_obs (Variable): The observation sample wrapped in Variable
    old_dist (dict): The dict of means and log_stds Variables of
        collected samples

    Returns
    -------
    kl_div (Variable): The KL-divergence between the old distribution
        and new distribution.
    """
    new_dist = policy.get_policy_distribution(all_obs)
    old_dist = policy.distribution(old_dist)

    kl_div = old_dist.kl_div(new_dist).mean()
    return kl_div