dreamer4 0.0.100__tar.gz → 0.0.101__tar.gz

{dreamer4-0.0.100 → dreamer4-0.0.101}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.0.100
+Version: 0.0.101
 Summary: Dreamer 4
 Project-URL: Homepage, https://pypi.org/project/dreamer4/
 Project-URL: Repository, https://github.com/lucidrains/dreamer4

{dreamer4-0.0.100 → dreamer4-0.0.101}/dreamer4/dreamer4.py

@@ -1900,6 +1900,7 @@ class DynamicsWorldModel(Module):
         gae_lambda = 0.95,
         ppo_eps_clip = 0.2,
         pmpo_pos_to_neg_weight = 0.5, # pos and neg equal weight
+        pmpo_reverse_kl = True,
         pmpo_kl_div_loss_weight = .3,
         value_clip = 0.4,
         policy_entropy_weight = .01,
@@ -2108,6 +2109,7 @@ class DynamicsWorldModel(Module):
 
         self.pmpo_pos_to_neg_weight = pmpo_pos_to_neg_weight
         self.pmpo_kl_div_loss_weight = pmpo_kl_div_loss_weight
+        self.pmpo_reverse_kl = pmpo_reverse_kl
 
         # rewards related
 
@@ -2578,11 +2580,18 @@ class DynamicsWorldModel(Module):
             # take care of kl
 
             if self.pmpo_kl_div_loss_weight > 0.:
+
                 new_unembedded_actions = self.action_embedder.unembed(policy_embed, pred_head_index = 0)
 
+                kl_div_inputs, kl_div_targets = new_unembedded_actions, old_action_unembeds
+
                 # mentioned that the "reverse direction for the prior KL" was used
+                # make optional, as observed instability in toy task
+
+                if self.pmpo_reverse_kl:
+                    kl_div_inputs, kl_div_targets = kl_div_targets, kl_div_inputs
 
-                discrete_kl_div, continuous_kl_div = self.action_embedder.kl_div(old_action_unembeds, new_unembedded_actions)
+                discrete_kl_div, continuous_kl_div = self.action_embedder.kl_div(kl_div_inputs, kl_div_targets)
 
                 # accumulate discrete and continuous kl div
 

{dreamer4-0.0.100 → dreamer4-0.0.101}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "dreamer4"
-version = "0.0.100"
+version = "0.0.101"
 description = "Dreamer 4"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }