Fix issue with KL calculation
@@ -177,7 +177,7 @@ if __name__ == "__main__":
logstd0_v = logstd_v.detach()
std_v = torch.exp(logstd_v)
std0_v = std_v.detach()
- kl = logstd_v - logstd0_v + (std0_v ** 2 + ((mu0_v - mu_v) ** 2) / (2.0 * std_v ** 2)) - 0.5
+ kl = logstd_v - logstd0_v + (std0_v ** 2 + (mu0_v - mu_v) ** 2) / (2.0 * std_v ** 2) - 0.5
return kl.sum(1, keepdim=True)
trpo.trpo_step(net_act, get_loss, get_kl, TRPO_MAX_KL, TRPO_DAMPING, device=device)
