metacontroller-pytorch 0.0.40__tar.gz → 0.0.42__tar.gz

{metacontroller_pytorch-0.0.40 → metacontroller_pytorch-0.0.42}/PKG-INFO

@@ -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.40 → metacontroller_pytorch-0.0.42}/README.md

@@ -1,6 +1,6 @@
 <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)
 
@@ -16,6 +16,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/metacontroller/__init__.py

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

{metacontroller_pytorch-0.0.40 → metacontroller_pytorch-0.0.42}/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 → metacontroller_pytorch-0.0.42}/pyproject.toml

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

{metacontroller_pytorch-0.0.40 → metacontroller_pytorch-0.0.42}/tests/test_metacontroller.py

@@ -7,7 +7,7 @@ from functools import partial
 
 import torch
 from torch import cat
-from metacontroller.metacontroller import Transformer, MetaController, policy_loss, z_score
+from metacontroller.metacontroller import Transformer, MetaController, policy_loss, z_score, extract_grpo_data
 from metacontroller.metacontroller_with_binary_mapper import MetaControllerWithBinaryMapper
 
 from memmap_replay_buffer import ReplayBuffer
@@ -109,10 +109,7 @@ def test_metacontroller(
         cache = None
         past_action_id = None
 
-        states = []
-        log_probs = []
-        switch_betas = []
-        latent_actions = []
+        grpo_data_list = []
 
         for one_state in subset_state.unbind(dim = 1):
             one_state = rearrange(one_state, 'b d -> b 1 d')
@@ -121,24 +118,19 @@ def test_metacontroller(
 
             past_action_id = model.action_readout.sample(logits)
 
-            # get log prob from meta controller latent actions
+            # extract grpo data and store
 
-            meta_output = cache.prev_hiddens.meta_controller
-
-            old_log_probs = meta_controller.log_prob(meta_output.action_dist, meta_output.actions)
-
-            states.append(meta_output.input_residual_stream)
-            log_probs.append(old_log_probs)
-            switch_betas.append(meta_output.switch_beta)
-            latent_actions.append(meta_output.actions)
+            grpo_data_list.append(extract_grpo_data(meta_controller, cache))
 
         # accumulate across time for the episode data
 
+        states, actions, log_probs, switch_betas = zip(*grpo_data_list)
+
         all_episodes.append((
             cat(states, dim = 1),
             cat(log_probs, dim = 1),
             cat(switch_betas, dim = 1),
-            cat(latent_actions, dim = 1)
+            cat(actions, dim = 1)
         ))
 
         all_rewards.append(torch.randn(1))

{metacontroller_pytorch-0.0.40 → metacontroller_pytorch-0.0.42}/train_babyai.py

@@ -26,7 +26,7 @@ from accelerate import Accelerator
 
 from babyai_env import create_env
 from memmap_replay_buffer import ReplayBuffer
-from metacontroller.metacontroller import Transformer, MetaController, policy_loss, z_score
+from metacontroller.metacontroller import Transformer, MetaController, policy_loss, z_score, extract_grpo_data
 from metacontroller.transformer_with_resnet import TransformerWithResnet
 
 # research entry point
@@ -195,13 +195,12 @@ def main(
 
                 # GRPO collection
 
-                meta_output = cache.prev_hiddens.meta_controller
-                old_log_probs = unwrapped_meta_controller.log_prob(meta_output.action_dist, meta_output.actions)
+                grpo_data = extract_grpo_data(unwrapped_meta_controller, cache)
 
-                states.append(meta_output.input_residual_stream)
-                log_probs.append(old_log_probs)
-                switch_betas.append(meta_output.switch_beta)
-                latent_actions.append(meta_output.actions)
+                states.append(grpo_data.state)
+                log_probs.append(grpo_data.log_prob)
+                switch_betas.append(grpo_data.switch_beta)
+                latent_actions.append(grpo_data.action)
 
                 next_state, reward, terminated, truncated, *_ = env.step(action.cpu().numpy())
 

metacontroller_pytorch-0.0.40/metacontroller/__init__.py

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