metacontroller-pytorch 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

metacontroller/metacontroller.py

@@ -22,7 +22,7 @@ from x_transformers import Decoder
 from x_mlps_pytorch import Feedforwards
 from x_evolution import EvoStrategy
 
-from discrete_continuous_embed_readout import Embed, Readout
+from discrete_continuous_embed_readout import Embed, Readout, EmbedAndReadout
 
 from assoc_scan import AssocScan
 
@@ -234,30 +234,25 @@ class Transformer(Module):
         self,
         dim,
         *,
-        embed: Embed | dict,
+        state_embed_readout: dict,
+        action_embed_readout: dict,
         lower_body: Decoder | dict,
         upper_body: Decoder | dict,
-        readout: Readout | dict,
         meta_controller: MetaController | None = None
     ):
         super().__init__()
 
-        if isinstance(embed, dict):
-            embed = Embed(dim = dim, **embed)
-
         if isinstance(lower_body, dict):
             lower_body = Decoder(dim = dim, **lower_body)
 
         if isinstance(upper_body, dict):
             upper_body = Decoder(dim = dim, **upper_body)
 
-        if isinstance(readout, dict):
-            readout = Readout(dim = dim, **readout)
+        self.state_embed, self.state_readout = EmbedAndReadout(dim, **state_embed_readout)
+        self.action_embed, self.action_readout = EmbedAndReadout(dim, **action_embed_readout)
 
-        self.embed = embed
         self.lower_body = lower_body
         self.upper_body = upper_body
-        self.readout = readout
 
         # meta controller
 
@@ -285,11 +280,13 @@ class Transformer(Module):
 
     def forward(
         self,
-        ids,
+        state,
+        action_ids,
         meta_controller: Module | None = None,
         cache: TransformerOutput | None = None,
         discovery_phase = False,
         meta_controller_temperature = 1.,
+        return_raw_action_dist = False,
         return_latents = False,
         return_cache = False,
     ):
@@ -297,21 +294,33 @@ class Transformer(Module):
 
         meta_controlling = exists(meta_controller)
 
+        behavioral_cloning = not meta_controlling and not return_raw_action_dist
+
         # by default, if meta controller is passed in, transformer is no grad
 
         lower_transformer_context = nullcontext if not meta_controlling else torch.no_grad
         meta_controller_context = nullcontext if meta_controlling else torch.no_grad
-        upper_transformer_context = nullcontext if meta_controlling and discovery_phase else torch.no_grad
+        upper_transformer_context = nullcontext if (not meta_controlling or discovery_phase) else torch.no_grad
 
         # handle cache
 
         lower_transformer_hiddens, meta_hiddens, upper_transformer_hiddens = cache.prev_hiddens if exists(cache) else ((None,) * 3)
 
+        # handle maybe behavioral cloning
+
+        if behavioral_cloning or (meta_controlling and discovery_phase):
+
+            state, target_state = state[:, :-1], state[:, 1:]
+            action_ids, target_action_ids = action_ids[:, :-1], action_ids[:, 1:]
+
         # transformer lower body
 
         with lower_transformer_context():
 
-            embed = self.embed(ids)
+            state_embed = self.state_embed(state)
+            action_embed = self.action_embed(action_ids)
+
+            embed = state_embed + action_embed
 
             residual_stream, next_lower_hiddens = self.lower_body(embed, cache = lower_transformer_hiddens, return_hiddens = True)
 
@@ -332,7 +341,23 @@ class Transformer(Module):
 
             # head readout
 
-            dist_params = self.readout(attended)
+            dist_params = self.action_readout(attended)
+
+        # maybe return behavior cloning loss
+
+        if behavioral_cloning:
+            state_dist_params = self.state_readout(attended)
+            state_clone_loss = self.state_readout.calculate_loss(state_dist_params, target_state)
+
+            action_clone_loss = self.action_readout.calculate_loss(dist_params, target_action_ids)
+
+            return state_clone_loss, action_clone_loss
+
+        elif meta_controlling and discovery_phase:
+
+            action_recon_loss = self.action_readout.calculate_loss(dist_params, target_action_ids)
+
+            return action_recon_loss, next_meta_hiddens.kl_loss
 
         # returning
 

{metacontroller_pytorch-0.0.14.dist-info → metacontroller_pytorch-0.0.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.14
+Version: 0.0.16
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -60,7 +60,7 @@ Implementation of the MetaController proposed in [Emergent temporal abstractions
 @misc{kobayashi2025emergenttemporalabstractionsautoregressive,
     title   = {Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning}, 
     author  = {Seijin Kobayashi and Yanick Schimpf and Maximilian Schlegel and Angelika Steger and Maciej Wolczyk and Johannes von Oswald and Nino Scherrer and Kaitlin Maile and Guillaume Lajoie and Blake A. Richards and Rif A. Saurous and James Manyika and Blaise Agüera y Arcas and Alexander Meulemans and João Sacramento},
-    year={2025},
+    year    = {2025},
     eprint  = {2512.20605},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},

metacontroller_pytorch-0.0.16.dist-info/RECORD

@@ -0,0 +1,6 @@
+metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
+metacontroller/metacontroller.py,sha256=BT7GH8F9NkEIYLEueBkkZ8glQ3Oht1FRoV84SIaTWdQ,11878
+metacontroller_pytorch-0.0.16.dist-info/METADATA,sha256=eyECb3994X58zyExLnnffMl3pOoMlIb-WAUhepIt0r8,3741
+metacontroller_pytorch-0.0.16.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+metacontroller_pytorch-0.0.16.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+metacontroller_pytorch-0.0.16.dist-info/RECORD,,

metacontroller_pytorch-0.0.14.dist-info/RECORD

@@ -1,6 +0,0 @@
-metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
-metacontroller/metacontroller.py,sha256=rQJLyXMHHNCZm0iohWqAkcMpYSi8b1Z0dgB-8AJVqqo,10751
-metacontroller_pytorch-0.0.14.dist-info/METADATA,sha256=-CP3Ak1NPaqTpyF4tTgwn-T47Pv2OiPzPFxecwGe3Ng,3736
-metacontroller_pytorch-0.0.14.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-metacontroller_pytorch-0.0.14.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-metacontroller_pytorch-0.0.14.dist-info/RECORD,,