Add KL divergence loss that mildly penalizes large updates in between time steps.
Add KL divergence loss that mildly penalizes large updates in between time steps.