From 4b06615018c0969be55041e4b2ee171c796b5dcf Mon Sep 17 00:00:00 2001 From: lucidrains Date: Wed, 29 Oct 2025 11:23:34 -0700 Subject: [PATCH] swap src and tgt for kl div --- dreamer4/dreamer4.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py index f3f260b..5c6f310 100644 --- a/dreamer4/dreamer4.py +++ b/dreamer4/dreamer4.py @@ -2580,7 +2580,7 @@ class DynamicsWorldModel(Module): if self.pmpo_kl_div_loss_weight > 0.: new_unembedded_actions = self.action_embedder.unembed(policy_embed, pred_head_index = 0) - discrete_kl_div, continuous_kl_div = self.action_embedder.kl_div(new_unembedded_actions, old_action_unembeds) + discrete_kl_div, continuous_kl_div = self.action_embedder.kl_div(old_action_unembeds, new_unembedded_actions) # accumulate discrete and continuous kl div diff --git a/pyproject.toml b/pyproject.toml index f583388..ea10f2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "dreamer4" -version = "0.0.96" +version = "0.0.97" description = "Dreamer 4" authors = [ { name = "Phil Wang", email = "lucidrains@gmail.com" }