metacontroller-pytorch 0.0.33__py3-none-any.whl → 0.0.34__py3-none-any.whl

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/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.dist-info → metacontroller_pytorch-0.0.34.dist-info}/METADATA

@@ -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.34.dist-info/RECORD

@@ -0,0 +1,8 @@
+metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
+metacontroller/metacontroller.py,sha256=Ii4Z2MuVMXJeWxAnZnRfSKrGv1t6_Y4R5BgvmtDrSW8,16203
+metacontroller/metacontroller_with_binary_mapper.py,sha256=PgXK7uk--gvZPk1h4WLdCBEA7m9Ji12xixa2wqLmsLY,9234
+metacontroller/transformer_with_resnet.py,sha256=R49ycusbq3kEX97WHZ41WY2ONc2mYPOuRUCmaFcBOEo,5546
+metacontroller_pytorch-0.0.34.dist-info/METADATA,sha256=xBFG34yRWTkcfJLFvC22jARLkR_W_6kcSsAV8r5UFWY,4747
+metacontroller_pytorch-0.0.34.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+metacontroller_pytorch-0.0.34.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+metacontroller_pytorch-0.0.34.dist-info/RECORD,,

metacontroller_pytorch-0.0.33.dist-info/RECORD

@@ -1,8 +0,0 @@
-metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
-metacontroller/metacontroller.py,sha256=B9XHYgVBrcJkEWhUORz--D5AHjcLnvLRVY9SRqVbhdw,16222
-metacontroller/metacontroller_with_binary_mapper.py,sha256=7vGtenScxvDQhkeYUmNnTTbVTJAtIFUVqEoWVGZP2Is,8936
-metacontroller/transformer_with_resnet.py,sha256=R49ycusbq3kEX97WHZ41WY2ONc2mYPOuRUCmaFcBOEo,5546
-metacontroller_pytorch-0.0.33.dist-info/METADATA,sha256=kWHDAnQeEueYWRzqPi5ouXjtSDIG2bnDqOfhznSUOoM,4747
-metacontroller_pytorch-0.0.33.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-metacontroller_pytorch-0.0.33.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-metacontroller_pytorch-0.0.33.dist-info/RECORD,,