metacontroller-pytorch 0.0.33__tar.gz → 0.0.34__tar.gz

{metacontroller_pytorch-0.0.33 → metacontroller_pytorch-0.0.34}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.33
+Version: 0.0.34
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller

{metacontroller_pytorch-0.0.33 → metacontroller_pytorch-0.0.34}/metacontroller/metacontroller.py

@@ -26,7 +26,7 @@ from discrete_continuous_embed_readout import Embed, Readout, EmbedAndReadout
 
 from assoc_scan import AssocScan
 
-from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, masked_mean, align_dims_left, pad_right_ndim_to
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, masked_mean, align_dims_left
 from torch_einops_utils.save_load import save_load
 
 # constants

{metacontroller_pytorch-0.0.33 → metacontroller_pytorch-0.0.34}/metacontroller/metacontroller_with_binary_mapper.py

@@ -23,7 +23,7 @@ from x_mlps_pytorch import Feedforwards
 
 from assoc_scan import AssocScan
 
-from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, masked_mean, align_dims_left, pad_right_ndim_to
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, align_dims_left
 from torch_einops_utils.save_load import save_load
 
 from vector_quantize_pytorch import BinaryMapper
@@ -143,22 +143,34 @@ class MetaControllerWithBinaryMapper(Module):
             *self.proposer_to_binary_logits.parameters()
         ]
 
+    def get_action_dist_for_internal_rl(
+        self,
+        residual_stream
+    ):
+        meta_embed = self.model_to_meta(residual_stream)
+
+        proposed_action_hidden, _ = self.action_proposer(meta_embed)
+
+        return self.proposer_to_binary_logits(proposed_action_hidden)
+
     def log_prob(
         self,
         action_dist,
         sampled_latent_action
     ):
-        action_prob = action_dist.sigmoid()
-        probs = stack((action_prob, 1. - action_prob), dim = -1)
-        log_probs = log(probs)
+        log_probs = stack((
+            F.logsigmoid(action_dist),
+            F.logsigmoid(-action_dist)
+        ), dim = -1)
 
         indices = sampled_latent_action.argmax(dim = -1)
         codes = self.binary_mapper.codes[indices].long()
 
         codes = rearrange(codes, '... -> ... 1')
         action_log_probs = log_probs.gather(-1, codes)
+        action_log_probs = rearrange(action_log_probs, '... 1 -> ...')
 
-        return rearrange(action_log_probs, '... 1 -> ...')
+        return action_log_probs.sum(dim = -1)
 
     def forward(
         self,

{metacontroller_pytorch-0.0.33 → metacontroller_pytorch-0.0.34}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "metacontroller-pytorch"
-version = "0.0.33"
+version = "0.0.34"
 description = "Transformer Metacontroller"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }