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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.41
+Version: 0.0.43
 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.43}/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.43}/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__(
@@ -273,7 +291,7 @@ class MetaController(Module):
         else:
             # else during inference, use the previous sampled latent action
 
-            assert seq_len == 1, f'inference RL phase must be done one token at a time'
+            assert seq_len == 1, 'inference RL phase must be done one token at a time - if replaying for policy optimization, please use `get_action_dist_for_internal_rl`'
             z_prev = prev_sampled_latent_action
 
         # switch input is previous latent action and the embedding

{metacontroller_pytorch-0.0.41 → metacontroller_pytorch-0.0.43}/metacontroller/metacontroller_with_binary_mapper.py

@@ -241,7 +241,7 @@ class MetaControllerWithBinaryMapper(Module):
         if discovery_phase:
             z_prev = cat((prev_sampled_code, sampled_codes[:, :-1]), dim = 1)
         else:
-            assert seq_len == 1, f'inference RL phase must be done one token at a time'
+            assert seq_len == 1, 'inference RL phase must be done one token at a time - if replaying for policy optimization, please use `get_action_dist_for_internal_rl`'
             z_prev = prev_sampled_code
 
         switch_input = torch.cat((meta_embed, z_prev), dim=-1)

{metacontroller_pytorch-0.0.41 → metacontroller_pytorch-0.0.43}/pyproject.toml

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

{metacontroller_pytorch-0.0.41 → metacontroller_pytorch-0.0.43}/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,36 +109,29 @@ 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')
 
-            logits, cache = model(one_state, past_action_id, meta_controller = meta_controller, return_cache = True)
+            logits, cache = model(one_state, past_action_id, meta_controller = meta_controller, cache = cache, return_cache = True)
 
             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 = extract_grpo_data(meta_controller, cache)
+            grpo_data_list.append(grpo_data)
 
         # 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))
@@ -153,6 +146,13 @@ def test_metacontroller(
     # simulate a policy loss update over the entire group
 
     group_states, group_log_probs, group_switch_betas, group_latent_actions = map(partial(cat, dim = 0), zip(*all_episodes))
+    
+    # parallel verification
+
+    parallel_action_dist = meta_controller.get_action_dist_for_internal_rl(group_states)
+    parallel_log_probs = meta_controller.log_prob(parallel_action_dist, group_latent_actions)
+
+    assert torch.allclose(parallel_log_probs, group_log_probs, atol = 1e-5), 'parallel log probs do not match stored log probs'
 
     for states, log_probs, switch_betas, latent_actions, advantages in zip(group_states, group_log_probs, group_switch_betas, group_latent_actions, group_advantages):
         replay_buffer.store_episode(

metacontroller_pytorch-0.0.43/train_baby_evo_strat.py

@@ -0,0 +1,213 @@
+# /// script
+# dependencies = [
+#   "fire",
+#   "gymnasium",
+#   "gymnasium[other]",
+#   "metacontroller-pytorch",
+#   "minigrid",
+#   "tqdm",
+#   "x-evolution",
+#   "einops"
+# ]
+# ///
+
+from __future__ import annotations
+import fire
+from pathlib import Path
+from shutil import rmtree
+import numpy as np
+
+import torch
+from torch import nn, Tensor, tensor
+from torch.nn import Module
+from einops import rearrange
+
+from babyai_env import create_env
+from metacontroller.metacontroller import Transformer, MetaController
+
+# functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# default fitness function
+
+def default_fitness_fn(
+    rewards: list[float],
+    states: list[any],
+    actions: list[any],
+    next_states: list[any],
+    infos: list[any]
+) -> float:
+    """
+    researchers can modify this function to engineer their own rewards and fitness scores
+    processing the entire episode at once for every noise vector of the population separately
+    """
+    return sum(rewards)
+
+# babyai environment for ES
+
+class BabyAIEnvironment(Module):
+    def __init__(
+        self,
+        env_id = 'BabyAI-BossLevel-v0',
+        video_folder = './recordings_babyai_es',
+        render_every_eps = 100,
+        max_steps = 500,
+        use_resnet = False,
+        fitness_fn = default_fitness_fn
+    ):
+        super().__init__()
+
+        self.env_id = env_id
+        self.video_folder = video_folder
+        self.render_every_eps = render_every_eps
+        self.max_steps = max_steps
+        self.use_resnet = use_resnet
+        self.fitness_fn = fitness_fn
+
+        # initial env creation for observation space etc. if needed
+        # but create_env is called inside pre_main_callback or reset
+        self.env = None
+
+    def pre_main_callback(self):
+        # clean up and initialize environment
+        rmtree(self.video_folder, ignore_errors = True)
+        
+        self.env = create_env(
+            self.env_id,
+            render_mode = 'rgb_array',
+            video_folder = self.video_folder,
+            render_every_eps = self.render_every_eps
+        )
+
+    def forward(self, model):
+        device = next(model.parameters()).device
+
+        seed = torch.randint(0, int(1e6), ()).item()
+        state, _ = self.env.reset(seed = seed)
+
+        step = 0
+        cache = None
+        past_action_id = None
+
+        unwrapped_model = getattr(model, 'model', model)
+
+        episode_rewards = []
+        episode_states = []
+        episode_actions = []
+        episode_next_states = []
+        episode_infos = []
+
+        while step < self.max_steps:
+            image = state['image']
+            image_tensor = torch.from_numpy(image).float().to(device)
+
+            if self.use_resnet:
+                image_tensor = rearrange(image_tensor, 'h w c -> 1 1 h w c')
+                image_tensor = unwrapped_model.visual_encode(image_tensor)
+            else:
+                image_tensor = rearrange(image_tensor, 'h w c -> 1 1 (h w c)')
+
+            if torch.is_tensor(past_action_id):
+                past_action_id = past_action_id.long()
+
+            with torch.no_grad():
+                logits, cache = model(
+                    image_tensor,
+                    past_action_id,
+                    return_cache = True,
+                    return_raw_action_dist = True,
+                    cache = cache
+                )
+
+            action = unwrapped_model.action_readout.sample(logits)
+            past_action_id = action
+            action_id = action.squeeze()
+
+            next_state, reward, terminated, truncated, info = self.env.step(action_id.cpu().numpy().item())
+
+            episode_rewards.append(reward)
+            episode_states.append(state)
+            episode_actions.append(action_id)
+            episode_next_states.append(next_state)
+            episode_infos.append(info)
+
+            done = terminated or truncated
+            if done:
+                break
+
+            state = next_state
+            step += 1
+
+        return self.fitness_fn(
+            episode_rewards,
+            episode_states,
+            episode_actions,
+            episode_next_states,
+            episode_infos
+        )
+
+def main(
+    env_id = 'BabyAI-BossLevel-v0',
+    num_generations = 100,
+    max_steps = 500,
+    render_every_eps = 100,
+    video_folder = './recordings_babyai_es',
+    transformer_weights_path: str | None = None,
+    meta_controller_weights_path: str | None = None,
+    output_meta_controller_path = 'metacontroller_es_trained.pt',
+    use_resnet = False,
+    noise_population_size = 50,
+    noise_scale = 1e-2,
+    learning_rate = 1e-3,
+    fitness_fn = default_fitness_fn
+):
+    # load model
+
+    assert exists(transformer_weights_path), "Transformer weights must be provided"
+    
+    # lazy import to avoid unnecessary dependencies if not used
+    from metacontroller.transformer_with_resnet import TransformerWithResnet as TransformerResnet
+    transformer_klass = TransformerResnet if use_resnet else Transformer
+    
+    model = transformer_klass.init_and_load(transformer_weights_path, strict = False)
+    model.eval()
+
+    if exists(meta_controller_weights_path):
+        meta_controller = MetaController.init_and_load(meta_controller_weights_path, strict = False)
+        model.meta_controller = meta_controller
+
+    assert exists(model.meta_controller), "MetaController must be present for evolution"
+
+    # setup environment
+    
+    babyai_env = BabyAIEnvironment(
+        env_id = env_id,
+        video_folder = video_folder,
+        render_every_eps = render_every_eps,
+        max_steps = max_steps,
+        use_resnet = use_resnet,
+        fitness_fn = fitness_fn
+    )
+
+    # evolve
+
+    model.evolve(
+        num_generations = num_generations,
+        environment = babyai_env,
+        noise_population_size = noise_population_size,
+        noise_scale = noise_scale,
+        learning_rate = learning_rate
+    )
+
+    # save
+
+    model.meta_controller.save(output_meta_controller_path)
+    print(f'MetaController weights saved to {output_meta_controller_path}')
+
+if __name__ == '__main__':
+    fire.Fire(main)

{metacontroller_pytorch-0.0.41 → metacontroller_pytorch-0.0.43}/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
@@ -57,22 +57,23 @@ def default(v, d):
 # main
 
 def main(
-    env_name: str = 'BabyAI-BossLevel-v0',
-    num_episodes: int = int(10e6),
-    max_timesteps: int = 500,
-    buffer_size: int = 5_000,
-    render_every_eps: int = 1_000,
-    video_folder: str = './recordings',
+    env_name = 'BabyAI-BossLevel-v0',
+    num_episodes = int(10e6),
+    max_timesteps = 500,
+    buffer_size = 5_000,
+    render_every_eps = 1_000,
+    video_folder = './recordings',
     seed: int | None = None,
     transformer_weights_path: str | None = None,
     meta_controller_weights_path: str | None = None,
-    output_meta_controller_path: str = 'metacontroller_rl_trained.pt',
-    use_resnet: bool = False,
-    lr: float = 1e-4,
-    num_groups: int = 16,
-    max_grad_norm: float = 1.0,
-    use_wandb: bool = False,
-    wandb_project: str = 'metacontroller-babyai-rl'
+    output_meta_controller_path = 'metacontroller_rl_trained.pt',
+    use_resnet = False,
+    lr = 1e-4,
+    batch_size = 16,
+    num_groups = 16,
+    max_grad_norm = 1.0,
+    use_wandb = False,
+    wandb_project = 'metacontroller-babyai-rl'
 ):
     # accelerator
 
@@ -195,13 +196,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())
 
@@ -264,7 +264,7 @@ def main(
         # learn
 
         if len(replay_buffer) >= buffer_size:
-            dl = replay_buffer.dataloader(batch_size = num_groups)
+            dl = replay_buffer.dataloader(batch_size = batch_size)
             dl = accelerator.prepare(dl)
 
             meta_controller.train()