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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.41
+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)
 

{metacontroller_pytorch-0.0.41 → 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)
 

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

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

{metacontroller_pytorch-0.0.41 → 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.41 → 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())