dreamer4 0.0.89__py3-none-any.whl → 0.0.91__py3-none-any.whl

dreamer4/dreamer4.py

@@ -186,6 +186,15 @@ def lens_to_mask(t, max_len = None):
 
     return einx.less('j, i -> i j', seq, t)
 
+def masked_mean(t, mask = None):
+    if not exists(mask):
+        return t.mean()
+
+    if not mask.any():
+        return t[mask].sum()
+
+    return t[mask].mean()
+
 def log(t, eps = 1e-20):
     return t.clamp(min = eps).log()
 
@@ -1824,6 +1833,7 @@ class DynamicsWorldModel(Module):
         gae_discount_factor = 0.997,
         gae_lambda = 0.95,
         ppo_eps_clip = 0.2,
+        pmpo_pos_to_neg_weight = 0.5, # pos and neg equal weight
         value_clip = 0.4,
         policy_entropy_weight = .01,
         gae_use_accelerated = False
@@ -2027,6 +2037,10 @@ class DynamicsWorldModel(Module):
         self.value_clip = value_clip
         self.policy_entropy_weight = value_clip
 
+        # pmpo related
+
+        self.pmpo_pos_to_neg_weight = pmpo_pos_to_neg_weight
+
         # rewards related
 
         self.keep_reward_ema_stats = keep_reward_ema_stats
@@ -2334,7 +2348,7 @@ class DynamicsWorldModel(Module):
         policy_optim: Optimizer | None = None,
         value_optim: Optimizer | None = None,
         only_learn_policy_value_heads = True, # in the paper, they do not finetune the entire dynamics model, they just learn the heads
-        use_signed_advantage = True,
+        use_pmpo = True,
         eps = 1e-6
     ):
 
@@ -2374,6 +2388,8 @@ class DynamicsWorldModel(Module):
         max_time = latents.shape[1]
         is_var_len = exists(experience.lens)
 
+        mask = None
+
         if is_var_len:
             learnable_lens = experience.lens - experience.is_truncated.long() # if is truncated, remove the last one, as it is bootstrapped value
             mask = lens_to_mask(learnable_lens, max_time)
@@ -2417,8 +2433,9 @@ class DynamicsWorldModel(Module):
         # apparently they just use the sign of the advantage
         # https://arxiv.org/abs/2410.04166v1
 
-        if use_signed_advantage:
-            advantage = advantage.sign()
+        if use_pmpo:
+            pos_advantage_mask = advantage >= 0.
+            neg_advantage_mask = ~pos_advantage_mask
         else:
             advantage = F.layer_norm(advantage, advantage.shape, eps = eps)
 
@@ -2464,35 +2481,48 @@ class DynamicsWorldModel(Module):
         log_probs = safe_cat(log_probs, dim = -1)
         entropies = safe_cat(entropies, dim = -1)
 
-        ratio = (log_probs - old_log_probs).exp()
+        advantage = rearrange(advantage, '... -> ... 1') # broadcast across all actions
 
-        clipped_ratio = ratio.clamp(1. - self.ppo_eps_clip, 1. + self.ppo_eps_clip)
+        if use_pmpo:
+            # pmpo - weighting the positive and negative advantages equally - ignoring magnitude of advantage and taking the sign
+            # seems to be weighted across batch and time, iiuc
+            # eq (10) in https://arxiv.org/html/2410.04166v1
 
-        advantage = rearrange(advantage, '... -> ... 1') # broadcast across all actions
+            if exists(mask):
+                pos_advantage_mask &= mask
+                neg_advantage_mask &= mask
+
+            α = self.pmpo_pos_to_neg_weight
 
-        # clipped surrogate loss
+            pos = masked_mean(log_probs, pos_advantage_mask)
+            neg = -masked_mean(log_probs, neg_advantage_mask)
 
-        policy_loss = -torch.min(ratio * advantage, clipped_ratio * advantage)
-        policy_loss = reduce(policy_loss, 'b t na -> b t', 'sum')
+            policy_loss = -(α * pos + (1. - α) * neg)
 
-        policy_loss = policy_loss.mean()
+        else:
+            # ppo clipped surrogate loss
+
+            ratio = (log_probs - old_log_probs).exp()
+            clipped_ratio = ratio.clamp(1. - self.ppo_eps_clip, 1. + self.ppo_eps_clip)
+
+            policy_loss = -torch.min(ratio * advantage, clipped_ratio * advantage)
+            policy_loss = reduce(policy_loss, 'b t na -> b t', 'sum')
+
+            policy_loss = masked_mean(policy_loss, mask)
 
         # handle entropy loss for naive exploration bonus
 
         entropy_loss = - reduce(entropies, 'b t na -> b t', 'sum')
 
+        entropy_loss = masked_mean(entropy_loss, mask)
+
+        # total policy loss
+
         total_policy_loss = (
             policy_loss +
             entropy_loss * self.policy_entropy_weight
         )
 
-        # maybe handle variable lengths
-
-        if is_var_len:
-            total_policy_loss = total_policy_loss[mask].mean()
-        else:
-            total_policy_loss = total_policy_loss.mean()
-
         # maybe take policy optimizer step
 
         if exists(policy_optim):
@@ -2543,6 +2573,7 @@ class DynamicsWorldModel(Module):
         batch_size = 1,
         agent_index = 0,
         tasks: int | Tensor | None = None,
+        latent_gene_ids = None,
         image_height = None,
         image_width = None,
         return_decoded_video = None,
@@ -2658,6 +2689,7 @@ class DynamicsWorldModel(Module):
                     step_sizes = step_size,
                     rewards = decoded_rewards,
                     tasks = tasks,
+                    latent_gene_ids = latent_gene_ids,
                     discrete_actions = decoded_discrete_actions,
                     continuous_actions = decoded_continuous_actions,
                     proprio = noised_proprio_with_context,

{dreamer4-0.0.89.dist-info → dreamer4-0.0.91.dist-info}/METADATA

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

dreamer4-0.0.91.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=LnCEoVek0ekyttB3z4axdgs_NszqpOg1LnNwhrfSisM,114742
+dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
+dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
+dreamer4-0.0.91.dist-info/METADATA,sha256=D2RUBJjn7zLpaRH0_duBFaAA0EMw5JuywscwecZUhEc,3065
+dreamer4-0.0.91.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.0.91.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.0.91.dist-info/RECORD,,

dreamer4-0.0.89.dist-info/RECORD

@@ -1,8 +0,0 @@
-dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
-dreamer4/dreamer4.py,sha256=F8UMxI8uAJmMHmXf_Xhibcs_PkMjNm4AW357U941REo,113725
-dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
-dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
-dreamer4-0.0.89.dist-info/METADATA,sha256=ak7JROeb_PRnOvAg8eoARtnbWh64-2JQJ5zVgSnwKpc,3065
-dreamer4-0.0.89.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dreamer4-0.0.89.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-dreamer4-0.0.89.dist-info/RECORD,,