RSSM Balanced KL Terms

Problem

Compute the two weighted KL terms of the Dreamer RSSM balance loss.

For univariate Gaussians, KL(post || prior) = ln(prior_std / post_std) + (post_std^2 + (post_mean - prior_mean)^2) / (2 * prior_std^2) - 0.5. Return [beta_dyn * KL, beta_rep * KL], the dynamics and representation shares of the loss. All standard deviations are positive.

Hint

Different weights on dynamics and representation gradient flow keep the prior from collapsing.

Test cases

case 1

in: [0,1,0,1,1,0.2]

expected: [0,0]

case 2

in: [0,1,1,1,1,0.2]

expected: [0.5,0.1]

case 3

in: [0,0.5,0,1,0.5,0.5]

expected: [0.15907359027997264,0.15907359027997264]

case 4

in: [2,0.5,1,2,1,0.1]

expected: [1.0425443611198906,0.10425443611198906]

Starter codePython

import math
def rssm_balanced_kl(post_mean, post_std, prior_mean, prior_std, beta_dyn, beta_rep):
    # Your code here
    pass

4 test cases · run your code against them in the editor