metacontroller-pytorch 0.0.40__py3-none-any.whl → 0.0.42__py3-none-any.whl

metacontroller/__init__.py

@@ -1 +1 @@
-from metacontroller.metacontroller import MetaController
+from metacontroller.metacontroller import MetaController, Transformer

metacontroller/metacontroller.py

@@ -66,6 +66,13 @@ MetaControllerOutput = namedtuple('MetaControllerOutput', (
     'switch_loss'
 ))
 
+GRPOOutput = namedtuple('GRPOOutput', (
+    'state',
+    'action',
+    'log_prob',
+    'switch_beta'
+))
+
 def z_score(t, eps = 1e-8):
     return (t - t.mean()) / (t.std() + eps)
 
@@ -107,6 +114,17 @@ def policy_loss(
 
     return masked_mean(losses, mask)
 
+def extract_grpo_data(meta_controller, transformer_output):
+    meta_output = transformer_output.prev_hiddens.meta_controller
+
+    state = meta_output.input_residual_stream
+    action = meta_output.actions
+    switch_beta = meta_output.switch_beta
+
+    log_prob = meta_controller.log_prob(meta_output.action_dist, action)
+
+    return GRPOOutput(state, action, log_prob, switch_beta)
+
 @save_load()
 class MetaController(Module):
     def __init__(

{metacontroller_pytorch-0.0.40.dist-info → metacontroller_pytorch-0.0.42.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.40
+Version: 0.0.42
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -53,7 +53,7 @@ Description-Content-Type: text/markdown
 
 <img src="./fig1.png" width="400px"></img>
 
-## metacontroller (wip)
+## metacontroller
 
 Implementation of the MetaController proposed in [Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning](https://arxiv.org/abs/2512.20605)
 
@@ -69,6 +69,91 @@ $ pip install metacontroller-pytorch
 
 - [Diego Calanzone](https://github.com/ddidacus) for proposing testing on BabyAI gridworld task, and submitting the [pull request](https://github.com/lucidrains/metacontroller/pull/3) for behavior cloning and discovery phase training for it!
 
+## Usage
+
+```python
+import torch
+from metacontroller import Transformer, MetaController
+
+# 1. initialize model
+
+model = Transformer(
+    dim = 512,
+    action_embed_readout = dict(num_discrete = 4),
+    state_embed_readout = dict(num_continuous = 384),
+    lower_body = dict(depth = 2),
+    upper_body = dict(depth = 2)
+)
+
+state = torch.randn(2, 128, 384)
+actions = torch.randint(0, 4, (2, 128))
+
+# 2. behavioral cloning (BC)
+
+state_loss, action_loss = model(state, actions)
+(state_loss + action_loss).backward()
+
+# 3. discovery phase
+
+meta_controller = MetaController(
+    dim_model = 512,
+    dim_meta_controller = 256,
+    dim_latent = 128
+)
+
+action_recon_loss, kl_loss, switch_loss = model(
+    state,
+    actions,
+    meta_controller = meta_controller,
+    discovery_phase = True
+)
+
+(action_recon_loss + kl_loss + switch_loss).backward()
+
+# 4. internal rl phase (GRPO)
+
+# ... collect trajectories ...
+
+logits, cache = model(
+    one_state,
+    past_action_id,
+    meta_controller = meta_controller,
+    return_cache = True
+)
+
+meta_output = cache.prev_hiddens.meta_controller
+old_log_probs = meta_controller.log_prob(meta_output.action_dist, meta_output.actions)
+
+# ... calculate advantages ...
+
+loss = meta_controller.policy_loss(
+    group_states,
+    group_old_log_probs,
+    group_latent_actions,
+    group_advantages,
+    group_switch_betas
+)
+
+loss.backward()
+```
+
+Or using [evolutionary strategies](https://arxiv.org/abs/2511.16652) for the last portion
+
+```python
+# 5. evolve (ES over GRPO)
+
+model.meta_controller = meta_controller
+
+def environment_callable(model):
+    # return a fitness score
+    return 1.0
+
+model.evolve(
+    num_generations = 10,
+    environment = environment_callable
+)
+```
+
 ## Citations
 
 ```bibtex

metacontroller_pytorch-0.0.42.dist-info/RECORD

@@ -0,0 +1,8 @@
+metacontroller/__init__.py,sha256=iSKbCDp3UrWhZg7SIJFYNjdVQU56u-vqZarE6qCSX74,70
+metacontroller/metacontroller.py,sha256=hOzMIeBwNZhrzpt6tnLahxuHJ4pPQ7JlEGBOxYHI_88,17875
+metacontroller/metacontroller_with_binary_mapper.py,sha256=Ce5-O95_pLuWNA3aZTlKrTGbc5cemb61tBtJBdSiLx4,9843
+metacontroller/transformer_with_resnet.py,sha256=R49ycusbq3kEX97WHZ41WY2ONc2mYPOuRUCmaFcBOEo,5546
+metacontroller_pytorch-0.0.42.dist-info/METADATA,sha256=f9KRrtFWHgZrx5HZYBGNrtfXrcfSOeZlRFfx7VYMOd0,6816
+metacontroller_pytorch-0.0.42.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+metacontroller_pytorch-0.0.42.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+metacontroller_pytorch-0.0.42.dist-info/RECORD,,

metacontroller_pytorch-0.0.40.dist-info/RECORD

@@ -1,8 +0,0 @@
-metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
-metacontroller/metacontroller.py,sha256=bhgCqqM-dfysGrMtZYe2w87lRVkf8fETjxUCdjrnI8Q,17386
-metacontroller/metacontroller_with_binary_mapper.py,sha256=Ce5-O95_pLuWNA3aZTlKrTGbc5cemb61tBtJBdSiLx4,9843
-metacontroller/transformer_with_resnet.py,sha256=R49ycusbq3kEX97WHZ41WY2ONc2mYPOuRUCmaFcBOEo,5546
-metacontroller_pytorch-0.0.40.dist-info/METADATA,sha256=0mvv1dI8pe3mNRLb8SJ3ZRWmBDvLgmz2wgI8j04xmQI,5114
-metacontroller_pytorch-0.0.40.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-metacontroller_pytorch-0.0.40.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-metacontroller_pytorch-0.0.40.dist-info/RECORD,,