metacontroller-pytorch 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

metacontroller/metacontroller.py

@@ -18,7 +18,7 @@ from einops.layers.torch import Rearrange
 
 # external modules
 
-from x_transformers import Decoder
+from x_transformers import Encoder, Decoder
 from x_mlps_pytorch import Feedforwards
 from x_evolution import EvoStrategy
 
@@ -72,7 +72,11 @@ class MetaController(Module):
         decoder_expansion_factor = 2.,
         decoder_depth = 1,
         hypernetwork_low_rank = 16,
-        assoc_scan_kwargs: dict = dict()
+        assoc_scan_kwargs: dict = dict(),
+        bidirectional_temporal_encoder_kwargs: dict = dict(
+            attn_dim_head = 32,
+            heads = 8
+        )
     ):
         super().__init__()
         dim_meta = default(dim_meta_controller, dim_model)
@@ -81,9 +85,9 @@ class MetaController(Module):
 
         self.model_to_meta = Linear(dim_model, dim_meta)
 
-        # there are two phases, the first (discovery ssl phase) uses acausal with some ssm i don't really believe in - let's just use a bidirectional GRU as placeholders
+        # there are two phases, the first (discovery ssl phase) uses acausal with some ssm i don't really believe in - let's just use bidirectional attention as placeholder
 
-        self.bidirectional_temporal_compressor = GRU(dim_meta, dim_meta, bidirectional = True) # revisit naming
+        self.bidirectional_temporal_encoder = Encoder(dim = dim_meta, depth = 1, **bidirectional_temporal_encoder_kwargs)
 
         self.emitter = GRU(dim_meta * 2, dim_meta * 2)
         self.emitter_to_action_mean_log_var = Readout(dim_meta * 2, num_continuous = dim_latent)
@@ -122,7 +126,7 @@ class MetaController(Module):
     def discovery_parameters(self):
         return [
             *self.model_to_meta.parameters(),
-            *self.bidirectional_temporal_compressor.parameters(),
+            *self.bidirectional_temporal_encoder.parameters(),
             *self.emitter.parameters(),
             *self.emitter_to_action_mean_log_var.parameters(),
             *self.decoder.parameters(),
@@ -157,10 +161,9 @@ class MetaController(Module):
         if discovery_phase:
             logger.warning('meta controller cache being passed back in for discovery phase, which does not make sense given bidirectional encoder')
 
-            temporal_compressed, _ = self.bidirectional_temporal_compressor(meta_embed)
-            temporal_compressed = reduce(temporal_compressed, '... (two d) -> ... d', 'mean', two = 2)
+            encoded_temporal = self.bidirectional_temporal_encoder(meta_embed)
 
-            proposed_action_hidden, _ = self.emitter(cat((temporal_compressed, meta_embed), dim = -1))
+            proposed_action_hidden, _ = self.emitter(cat((encoded_temporal, meta_embed), dim = -1))
             readout = self.emitter_to_action_mean_log_var
 
         else: # else internal rl phase
@@ -256,7 +259,7 @@ class Transformer(Module):
         super().__init__()
 
         if isinstance(lower_body, dict):
-            lower_body = Decoder(dim = dim, **lower_body)
+            lower_body = Decoder(dim = dim, pre_norm_has_final_norm = False, **lower_body)
 
         if isinstance(upper_body, dict):
             upper_body = Decoder(dim = dim, **upper_body)

{metacontroller_pytorch-0.0.18.dist-info → metacontroller_pytorch-0.0.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.18
+Version: 0.0.20
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -78,3 +78,15 @@ Implementation of the MetaController proposed in [Emergent temporal abstractions
     url     = {https://api.semanticscholar.org/CorpusID:279464702}
 }
 ```
+
+```bibtex
+@misc{fleuret2025freetransformer,
+    title     = {The Free Transformer}, 
+    author    = {François Fleuret},
+    year      = {2025},
+    eprint    = {2510.17558},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url       = {https://arxiv.org/abs/2510.17558}, 
+}
+```

metacontroller_pytorch-0.0.20.dist-info/RECORD

@@ -0,0 +1,6 @@
+metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
+metacontroller/metacontroller.py,sha256=3QZrId9z8I6MMQ3GhEQ6Xb5LFRTFJq4EAU4JCvRmm-4,12368
+metacontroller_pytorch-0.0.20.dist-info/METADATA,sha256=5t4rDJiJzbx7m9BNsTTgO5JOnavaX-3jv31HTGuLP6A,4034
+metacontroller_pytorch-0.0.20.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+metacontroller_pytorch-0.0.20.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+metacontroller_pytorch-0.0.20.dist-info/RECORD,,

metacontroller_pytorch-0.0.18.dist-info/RECORD

@@ -1,6 +0,0 @@
-metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
-metacontroller/metacontroller.py,sha256=xrsP8YyYFZ_Z4rZx0BiYJZT2Q3zzXZppZnPKZfg-mtg,12306
-metacontroller_pytorch-0.0.18.dist-info/METADATA,sha256=fwk9OgNmoPBbZYsezBTerlcAd6iPMpK2zWXFPDlhrs4,3741
-metacontroller_pytorch-0.0.18.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-metacontroller_pytorch-0.0.18.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-metacontroller_pytorch-0.0.18.dist-info/RECORD,,