metacontroller-pytorch 0.0.31__tar.gz → 0.0.32__tar.gz

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.32}/PKG-INFO

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

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.32}/metacontroller/metacontroller.py

@@ -107,7 +107,6 @@ class MetaController(Module):
 
         self.switch_per_latent_dim = switch_per_latent_dim
 
-
         self.dim_latent = dim_latent
         self.switching_unit = GRU(dim_meta + dim_latent, dim_meta)
         self.to_switching_unit_beta = nn.Linear(dim_meta, dim_latent if switch_per_latent_dim else 1, bias = False)
@@ -147,6 +146,13 @@ class MetaController(Module):
             *self.action_proposer_mean_log_var.parameters()
         ]
 
+    def log_prob(
+        self,
+        action_dist,
+        sampled_latent_action
+    ):
+        return self.action_proposer_mean_log_var.log_prob(action_dist, sampled_latent_action)
+
     def forward(
         self,
         residual_stream,
@@ -276,6 +282,12 @@ class MetaController(Module):
 
 # main transformer, which is subsumed into the environment after behavioral cloning
 
+Hiddens = namedtuple('Hiddens', (
+    'lower_body',
+    'meta_controller',
+    'upper_body'
+))
+
 TransformerOutput = namedtuple('TransformerOutput', (
     'residual_stream_latent',
     'prev_hiddens'
@@ -441,4 +453,4 @@ class Transformer(Module):
         if return_one:
             return dist_params
 
-        return dist_params, TransformerOutput(residual_stream, (next_lower_hiddens, next_meta_hiddens, next_upper_hiddens))
+        return dist_params, TransformerOutput(residual_stream, Hiddens(next_lower_hiddens, next_meta_hiddens, next_upper_hiddens))

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.32}/metacontroller/metacontroller_with_binary_mapper.py

@@ -50,6 +50,9 @@ def default(*args):
 def straight_through(src, tgt):
     return tgt + src - src.detach()
 
+def log(t, eps = 1e-20):
+    return t.clamp_min(eps).log()
+
 # meta controller
 
 @save_load()
@@ -137,6 +140,23 @@ class MetaControllerWithBinaryMapper(Module):
             *self.proposer_to_binary_logits.parameters()
         ]
 
+    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)
+
+        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)
+
+        return rearrange(action_log_probs, '... 1 -> ...')
+
     def forward(
         self,
         residual_stream,

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.32}/pyproject.toml

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

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.32}/tests/test_metacontroller.py

@@ -80,6 +80,12 @@ def test_metacontroller(
         assert logits.shape == (2, 1, *assert_shape)
         past_action_id = model.action_readout.sample(logits)
 
+        # get log prob from meta controller latent actions
+
+        meta_controller_hidden = cache.prev_hiddens.meta_controller
+
+        old_log_probs = meta_controller.log_prob(meta_controller_hidden.action_dist, meta_controller_hidden.actions)
+
     # evolutionary strategies over grpo
 
     model.meta_controller = meta_controller