metacontroller-pytorch 0.0.36__tar.gz → 0.0.38__tar.gz

{metacontroller_pytorch-0.0.36 → metacontroller_pytorch-0.0.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.36
+Version: 0.0.38
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -94,6 +94,18 @@ $ pip install metacontroller-pytorch
 }
 ```
 
+```bibtex
+@misc{hwang2025dynamicchunkingendtoendhierarchical,
+    title   = {Dynamic Chunking for End-to-End Hierarchical Sequence Modeling},
+    author  = {Sukjun Hwang and Brandon Wang and Albert Gu},
+    year    = {2025},
+    eprint  = {2507.07955},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2507.07955},
+}
+```
+
 ```bibtex
 @misc{fleuret2025freetransformer,
     title     = {The Free Transformer}, 

{metacontroller_pytorch-0.0.36 → metacontroller_pytorch-0.0.38}/README.md

@@ -41,6 +41,18 @@ $ pip install metacontroller-pytorch
 }
 ```
 
+```bibtex
+@misc{hwang2025dynamicchunkingendtoendhierarchical,
+    title   = {Dynamic Chunking for End-to-End Hierarchical Sequence Modeling},
+    author  = {Sukjun Hwang and Brandon Wang and Albert Gu},
+    year    = {2025},
+    eprint  = {2507.07955},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2507.07955},
+}
+```
+
 ```bibtex
 @misc{fleuret2025freetransformer,
     title     = {The Free Transformer}, 

{metacontroller_pytorch-0.0.36 → metacontroller_pytorch-0.0.38}/metacontroller/metacontroller.py

@@ -126,6 +126,7 @@ class MetaController(Module):
         )
     ):
         super().__init__()
+        self.dim_model = dim_model
         dim_meta = default(dim_meta_controller, dim_model)
 
         # the linear that brings from model dimension 
@@ -171,6 +172,15 @@ class MetaController(Module):
 
         self.register_buffer('zero', tensor(0.), persistent = False)
 
+    @property
+    def replay_buffer_field_dict(self):
+        return dict(
+            states = ('float', self.dim_model),
+            log_probs = ('float', self.dim_latent),
+            switch_betas = ('float', self.dim_latent if self.switch_per_latent_dim else 1),
+            latent_actions = ('float', self.dim_latent)
+        )
+
     def discovery_parameters(self):
         return [
             *self.model_to_meta.parameters(),
@@ -408,6 +418,7 @@ class Transformer(Module):
         meta_controller: Module | None = None,
         cache: TransformerOutput | None = None,
         discovery_phase = False,
+        force_behavior_cloning = False,
         meta_controller_temperature = 1.,
         return_raw_action_dist = False,
         return_latents = False,
@@ -420,11 +431,15 @@ class Transformer(Module):
 
         meta_controller = default(meta_controller, self.meta_controller)
 
+        if force_behavior_cloning:
+            assert not discovery_phase, 'discovery phase cannot be set to True if force behavioral cloning is set to True'
+            meta_controller = None
+
         has_meta_controller = exists(meta_controller)
 
         assert not (discovery_phase and not has_meta_controller), 'meta controller must be made available during discovery phase'
 
-        behavioral_cloning = not has_meta_controller and not return_raw_action_dist
+        behavioral_cloning = force_behavior_cloning or (not has_meta_controller and not return_raw_action_dist)
 
         # by default, if meta controller is passed in, transformer is no grad
 
@@ -472,7 +487,7 @@ class Transformer(Module):
 
         with meta_controller_context():
 
-            if exists(meta_controller):
+            if exists(meta_controller) and not behavioral_cloning:
                 control_signal, next_meta_hiddens = meta_controller(residual_stream, cache = meta_hiddens, discovery_phase = discovery_phase, temperature = meta_controller_temperature, episode_lens = episode_lens)
             else:
                 control_signal, next_meta_hiddens = self.zero, None

{metacontroller_pytorch-0.0.36 → metacontroller_pytorch-0.0.38}/metacontroller/metacontroller_with_binary_mapper.py

@@ -74,7 +74,7 @@ class MetaControllerWithBinaryMapper(Module):
         kl_loss_threshold = 0.
     ):
         super().__init__()
-
+        self.dim_model = dim_model
         assert not switch_per_code, 'switch_per_code is not supported for binary mapper'
 
         dim_meta = default(dim_meta_controller, dim_model)
@@ -126,6 +126,15 @@ class MetaControllerWithBinaryMapper(Module):
 
         self.register_buffer('zero', tensor(0.), persistent = False)
 
+    @property
+    def replay_buffer_field_dict(self):
+        return dict(
+            states = ('float', self.dim_model),
+            log_probs = ('float', self.dim_code_bits),
+            switch_betas = ('float', self.num_codes if self.switch_per_code else 1),
+            latent_actions = ('float', self.num_codes)
+        )
+
     def discovery_parameters(self):
         return [
             *self.model_to_meta.parameters(),

{metacontroller_pytorch-0.0.36 → metacontroller_pytorch-0.0.38}/pyproject.toml

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

{metacontroller_pytorch-0.0.36 → metacontroller_pytorch-0.0.38}/tests/test_metacontroller.py

@@ -69,12 +69,6 @@ def test_metacontroller(
             dim_latent = 128,
             switch_per_latent_dim = switch_per_latent_dim
         )
-
-        field_shapes = dict(
-            log_probs = ('float', 128),
-            switch_betas = ('float', 128 if switch_per_latent_dim else 1),
-            latent_actions = ('float', 128)
-        )
     else:
         meta_controller = MetaControllerWithBinaryMapper(
             dim_model = 512,
@@ -83,12 +77,6 @@ def test_metacontroller(
             dim_code_bits = 8, # 2 ** 8 = 256 codes
         )
 
-        field_shapes = dict(
-            log_probs = ('float', 8),
-            switch_betas = ('float', 8 if switch_per_latent_dim else 1),
-            latent_actions = ('float', 256)
-        )
-
     # discovery phase
 
     (action_recon_loss, kl_loss, switch_loss) = model(state, actions, meta_controller = meta_controller, discovery_phase = True, episode_lens = episode_lens)
@@ -104,10 +92,7 @@ def test_metacontroller(
         test_folder,
         max_episodes = 3,
         max_timesteps = 256,
-        fields = dict(
-            states = ('float', 512),
-            **field_shapes
-        ),
+        fields = meta_controller.replay_buffer_field_dict,
         meta_fields = dict(
             advantages = 'float'
         )

{metacontroller_pytorch-0.0.36 → metacontroller_pytorch-0.0.38}/train_behavior_clone_babyai.py

@@ -130,8 +130,8 @@ def train(
     for epoch in range(cloning_epochs + discovery_epochs):
 
         model.train()
-        total_state_loss = 0.
-        total_action_loss = 0.
+        from collections import defaultdict
+        total_losses = defaultdict(float)
 
         progress_bar = tqdm(dataloader, desc = f"Epoch {epoch}", disable = not accelerator.is_local_main_process)
 
@@ -200,9 +200,9 @@ def train(
                 optim.zero_grad()
 
             # log
-
-            total_state_loss += state_loss.item()
-            total_action_loss += action_loss.item()
+            
+            for key, value in log.items():
+                total_losses[key] += value
 
             accelerator.log({
                 **log,
@@ -210,15 +210,11 @@ def train(
                 "grad_norm": grad_norm.item()
             })
 
-            progress_bar.set_postfix(
-                state_loss = state_loss.item(),
-                action_loss = action_loss.item()
-            )
-
-        avg_state_loss = total_state_loss / len(dataloader)
-        avg_action_loss = total_action_loss / len(dataloader)
+            progress_bar.set_postfix(**log)
 
-        accelerator.print(f"Epoch {epoch}: state_loss={avg_state_loss:.4f}, action_loss={avg_action_loss:.4f}")
+        avg_losses = {k: v / len(dataloader) for k, v in total_losses.items()}
+        avg_losses_str = ", ".join([f"{k}={v:.4f}" for k, v in avg_losses.items()])
+        accelerator.print(f"Epoch {epoch}: {avg_losses_str}")
 
     # save weights
 
@@ -231,7 +227,7 @@ def train(
         unwrapped_meta_controller = accelerator.unwrap_model(meta_controller)
         unwrapped_meta_controller.save(meta_controller_checkpoint_path)
 
-        accelerator.print(f"Model saved to {checkpoint_path}, MetaControler to {meta_controller_checkpoint_path}")
+        accelerator.print(f"Model saved to {checkpoint_path}, MetaController to {meta_controller_checkpoint_path}")
 
     accelerator.end_training()