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

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
@@ -198,6 +198,14 @@ def masked_mean(t, mask = None):
 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)]
 
@@ -824,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:
@@ -842,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,
         *,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreamer4
-Version: 0.0.91
+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.92.dist-info/RECORD

@@ -0,0 +1,8 @@
+dreamer4/__init__.py,sha256=Jssh1obzDRtTfBLZl36kXge1cIQlMjf_8DyjPulvKSk,183
+dreamer4/dreamer4.py,sha256=8RIgxdlgphe19w4bdcy2w7qWqVL4CZtoPfL6LXwj7jE,116380
+dreamer4/mocks.py,sha256=TfqOB_Gq6N_GggBYwa6ZAJQx38ntlYbXZe23Ne4jshw,2502
+dreamer4/trainers.py,sha256=D2b7WTgTHElLhIWLFgl2Ct2knGJLTk91HHpC5UkNvG0,14028
+dreamer4-0.0.92.dist-info/METADATA,sha256=we70JHu-GvXjlcb5HR2bgtn8dc6L-uOARiNdqAKqDTU,3065
+dreamer4-0.0.92.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dreamer4-0.0.92.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+dreamer4-0.0.92.dist-info/RECORD,,

dreamer4-0.0.91.dist-info/RECORD

@@ -1,8 +0,0 @@
-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,,