dreamer4 0.0.90__tar.gz → 0.0.92__tar.gz

{dreamer4-0.0.90 → dreamer4-0.0.92}/PKG-INFO

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

{dreamer4-0.0.90 → dreamer4-0.0.92}/dreamer4/dreamer4.py

@@ -11,7 +11,7 @@ from dataclasses import dataclass, asdict
 import torch
 import torch.nn.functional as F
 from torch.nested import nested_tensor
-from torch.distributions import Normal
+from torch.distributions import Normal, kl
 from torch.nn import Module, ModuleList, Embedding, Parameter, Sequential, Linear, RMSNorm, Identity
 from torch import nn, cat, stack, arange, tensor, Tensor, is_tensor, full, zeros, ones, randint, rand, randn, randn_like, empty, full, linspace, arange
 from torch.utils._pytree import tree_flatten, tree_unflatten
@@ -186,9 +186,26 @@ 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()
 
+def mean_log_var_to_distr(
+    mean_log_var: Tensor
+) -> Normal:
+
+    mean, log_var = mean_log_var.unbind(dim = -1)
+    std = (0.5 * log_var).exp()
+    return Normal(mean, std)
+
 def safe_cat(tensors, dim):
     tensors = [*filter(exists, tensors)]
 
@@ -815,10 +832,7 @@ class ActionEmbedder(Module):
         continuous_entropies = None
 
         if exists(continuous_targets):
-            mean, log_var = continuous_action_mean_log_var.unbind(dim = -1)
-            std = (0.5 * log_var).exp()
-
-            distr = Normal(mean, std)
+            distr = mean_log_var_to_distr(continuous_action_mean_log_var)
             continuous_log_probs = distr.log_prob(continuous_targets)
 
             if return_entropies:
@@ -833,6 +847,54 @@ class ActionEmbedder(Module):
 
         return log_probs, entropies
 
+    def kl_div(
+        self,
+        src: tuple[Tensor | None, Tensor | None],
+        tgt: tuple[Tensor | None, Tensor | None]
+    ) -> tuple[Tensor | None, Tensor | None]:
+
+        src_discrete, src_continuous = src
+        tgt_discrete, tgt_continuous = tgt
+
+        discrete_kl_div = None
+
+        # split discrete if it is not already (multiple discrete actions)
+
+        if exists(src_discrete):
+
+            discrete_split = self.num_discrete_actions.tolist()
+
+            if is_tensor(src_discrete):
+                src_discrete = src_discrete.split(discrete_split, dim = -1)
+
+            if is_tensor(tgt_discrete):
+                tgt_discrete = tgt_discrete.split(discrete_split, dim = -1)
+
+            discrete_kl_divs = []
+
+            for src_logit, tgt_logit in zip(src_discrete, tgt_discrete):
+
+                src_log_probs = src_logit.log_softmax(dim = -1)
+                tgt_prob = tgt_logit.softmax(dim = -1)
+
+                one_discrete_kl_div = F.kl_div(src_log_probs, tgt_prob, reduction = 'none')
+
+                discrete_kl_divs.append(one_discrete_kl_div.sum(dim = -1))
+
+            discrete_kl_div = stack(discrete_kl_divs, dim = -1)
+
+        # calculate kl divergence for continuous
+
+        continuous_kl_div = None
+
+        if exists(src_continuous):
+            src_normal = mean_log_var_to_distr(src_continuous)
+            tgt_normal = mean_log_var_to_distr(tgt_continuous)
+
+            continuous_kl_div = kl.kl_divergence(src_normal, tgt_normal)
+
+        return discrete_kl_div, continuous_kl_div
+
     def forward(
         self,
         *,
@@ -1824,6 +1886,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 +2090,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 +2401,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 +2441,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 +2486,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 +2534,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
 
-        # clipped surrogate loss
+            α = self.pmpo_pos_to_neg_weight
 
-        policy_loss = -torch.min(ratio * advantage, clipped_ratio * advantage)
-        policy_loss = reduce(policy_loss, 'b t na -> b t', 'sum')
+            pos = masked_mean(log_probs, pos_advantage_mask)
+            neg = -masked_mean(log_probs, neg_advantage_mask)
 
-        policy_loss = policy_loss.mean()
+            policy_loss = -(α * pos + (1. - α) * neg)
+
+        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):

{dreamer4-0.0.90 → dreamer4-0.0.92}/pyproject.toml

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

{dreamer4-0.0.90 → dreamer4-0.0.92}/tests/test_dreamer.py

@@ -346,6 +346,15 @@ def test_action_embedder():
     assert discrete_logits.shape == (2, 3, 8)
     assert continuous_mean_log_var.shape == (2, 3, 2, 2)
 
+    # test kl div
+
+    discrete_logits_tgt, continuous_mean_log_var_tgt = embedder.unembed(action_embed)
+
+    discrete_kl_div, continuous_kl_div = embedder.kl_div((discrete_logits, continuous_mean_log_var), (discrete_logits_tgt, continuous_mean_log_var_tgt))
+
+    assert discrete_kl_div.shape == (2, 3, 2)
+    assert continuous_kl_div.shape == (2, 3, 2)
+
     # return discrete split by number of actions
 
     discrete_logits, continuous_mean_log_var = embedder.unembed(action_embed, return_split_discrete = True)
@@ -611,13 +620,13 @@ def test_cache_generate():
     generated, time_kv_cache = dynamics.generate(1, time_kv_cache = time_kv_cache, return_time_kv_cache = True)
 
 @param('vectorized', (False, True))
-@param('use_signed_advantage', (False, True))
+@param('use_pmpo', (False, True))
 @param('env_can_terminate', (False, True))
 @param('env_can_truncate', (False, True))
 @param('store_agent_embed', (False, True))
 def test_online_rl(
     vectorized,
-    use_signed_advantage,
+    use_pmpo,
     env_can_terminate,
     env_can_truncate,
     store_agent_embed
@@ -674,7 +683,7 @@ def test_online_rl(
     if store_agent_embed:
         assert exists(combined_experience.agent_embed)
 
-    actor_loss, critic_loss = world_model_and_policy.learn_from_experience(combined_experience, use_signed_advantage = use_signed_advantage)
+    actor_loss, critic_loss = world_model_and_policy.learn_from_experience(combined_experience, use_pmpo = use_pmpo)
 
     actor_loss.backward()
     critic_loss.backward()