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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.37
+Version: 0.0.40
 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.37 → metacontroller_pytorch-0.0.40}/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.40/babyai_env.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+from shutil import rmtree
+
+import gymnasium as gym
+import minigrid
+from minigrid.wrappers import FullyObsWrapper, SymbolicObsWrapper
+
+# functions
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+# env creation
+
+def create_env(
+    env_id,
+    render_mode = 'rgb_array',
+    video_folder = None,
+    render_every_eps = 1000
+):
+    # register minigrid environments if needed
+    minigrid.register_minigrid_envs()
+
+    # environment
+    env = gym.make(env_id, render_mode = render_mode)
+    env = FullyObsWrapper(env)
+    env = SymbolicObsWrapper(env)
+
+    if video_folder is not None:
+        video_folder = Path(video_folder)
+        rmtree(video_folder, ignore_errors = True)
+
+        env = gym.wrappers.RecordVideo(
+            env = env,
+            video_folder = str(video_folder),
+            name_prefix = 'babyai',
+            episode_trigger = lambda eps_num: divisible_by(eps_num, render_every_eps),
+            disable_logger = True
+        )
+
+    return env

{metacontroller_pytorch-0.0.37 → metacontroller_pytorch-0.0.40}/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(),

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

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

metacontroller_pytorch-0.0.40/test_babyai_e2e.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+set -e
+
+# 1. Gather trajectories
+echo "Gathering trajectories..."
+uv run gather_babyai_trajs.py \
+    --num_seeds 100 \
+    --num_episodes_per_seed 10 \
+    --num_steps 500 \
+    --output_dir end_to_end_trajectories \
+    --env_id BabyAI-MiniBossLevel-v0
+
+# 2. Behavioral cloning
+echo "Training behavioral cloning model..."
+ACCELERATE_USE_CPU=true ACCELERATE_MIXED_PRECISION=no uv run train_behavior_clone_babyai.py \
+    --cloning_epochs 10 \
+    --discovery_epochs 10 \
+    --batch_size 256 \
+    --input_dir end_to_end_trajectories \
+    --env_id BabyAI-MiniBossLevel-v0 \
+    --checkpoint_path end_to_end_model.pt \
+    --use_resnet
+
+# 3. Inference rollouts
+echo "Running inference rollouts..."
+uv run train_babyai.py \
+    --transformer_weights_path end_to_end_model.pt \
+    --meta_controller_weights_path meta_controller_discovery.pt \
+    --env_name BabyAI-MiniBossLevel-v0 \
+    --num_episodes 1000 \
+    --buffer_size 1000 \
+    --max_timesteps 100 \
+    --num_groups 16 \
+    --lr 1e-4 \
+    --use_resnet

{metacontroller_pytorch-0.0.37 → metacontroller_pytorch-0.0.40}/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.40/train_babyai.py

@@ -0,0 +1,314 @@
+# /// script
+# dependencies = [
+#   "fire",
+#   "gymnasium",
+#   "gymnasium[other]",
+#   "memmap-replay-buffer>=0.0.12",
+#   "metacontroller-pytorch",
+#   "minigrid",
+#   "tqdm"
+# ]
+# ///
+
+from fire import Fire
+from pathlib import Path
+from functools import partial
+from shutil import rmtree
+from tqdm import tqdm
+
+import torch
+from torch import cat, tensor, stack
+from torch.optim import Adam
+
+from einops import rearrange
+
+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.transformer_with_resnet import TransformerWithResnet
+
+# research entry point
+
+def reward_shaping_fn(
+    cumulative_rewards: torch.Tensor,
+    all_rewards: torch.Tensor,
+    episode_lens: torch.Tensor
+) -> torch.Tensor | None:
+    """
+    researchers can modify this function to engineer rewards
+    or return None to reject the entire batch
+    
+    cumulative_rewards: (num_episodes,)
+    all_rewards: (num_episodes, max_timesteps)
+    episode_lens: (num_episodes,)
+    """
+    return cumulative_rewards
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else 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',
+    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'
+):
+    # accelerator
+
+    accelerator = Accelerator(log_with = 'wandb' if use_wandb else None)
+
+    if use_wandb:
+        accelerator.init_trackers(wandb_project)
+
+    # environment
+
+    env = create_env(
+        env_name,
+        render_mode = 'rgb_array',
+        video_folder = video_folder,
+        render_every_eps = render_every_eps
+    )
+
+    # load models
+
+    model = None
+    if exists(transformer_weights_path):
+        weights_path = Path(transformer_weights_path)
+        assert weights_path.exists(), f"transformer weights not found at {weights_path}"
+        
+        transformer_klass = TransformerWithResnet if use_resnet else Transformer
+        model = transformer_klass.init_and_load(str(weights_path), strict = False)
+        model.eval()
+
+    meta_controller = None
+    if exists(meta_controller_weights_path):
+        weights_path = Path(meta_controller_weights_path)
+        assert weights_path.exists(), f"meta controller weights not found at {weights_path}"
+        meta_controller = MetaController.init_and_load(str(weights_path), strict = False)
+        meta_controller.eval()
+
+    meta_controller = default(meta_controller, getattr(model, 'meta_controller', None))
+    assert exists(meta_controller), "MetaController must be present for reinforcement learning"
+
+    # optimizer
+
+    optim = Adam(meta_controller.internal_rl_parameters(), lr = lr)
+
+    # prepare
+
+    model, meta_controller, optim = accelerator.prepare(model, meta_controller, optim)
+
+    unwrapped_model = accelerator.unwrap_model(model)
+    unwrapped_meta_controller = accelerator.unwrap_model(meta_controller)
+
+    # replay buffer
+
+    replay_buffer = ReplayBuffer(
+        './replay-data',
+        max_episodes = buffer_size,
+        max_timesteps = max_timesteps + 1,
+        fields = meta_controller.replay_buffer_field_dict,
+        meta_fields = dict(advantages = 'float'),
+        overwrite = True,
+        circular = True
+    )
+
+    # rollouts
+
+    num_batch_updates = num_episodes // num_groups
+
+    pbar = tqdm(range(num_batch_updates), desc = 'training')
+
+    for _ in pbar:
+
+        all_episodes = []
+        all_cumulative_rewards = []
+        all_step_rewards = []
+        all_episode_lens = []
+
+        group_seed = default(seed, torch.randint(0, 1000000, (1,)).item())
+
+        for _ in range(num_groups):
+
+            state, *_ = env.reset(seed = group_seed)
+
+            cache = None
+            past_action_id = None
+
+            states = []
+            log_probs = []
+            switch_betas = []
+            latent_actions = []
+
+            total_reward = 0.
+            step_rewards = []
+            episode_len = max_timesteps
+
+            for step in range(max_timesteps):
+
+                image = state['image']
+                image_tensor = torch.from_numpy(image).float().to(accelerator.device)
+
+                if use_resnet:
+                    image_tensor = rearrange(image_tensor, 'h w c -> 1 1 h w c')
+                    image_tensor = 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 = unwrapped_model(
+                        image_tensor,
+                        past_action_id,
+                        meta_controller = unwrapped_meta_controller,
+                        return_cache = True,
+                        return_raw_action_dist = True,
+                        cache = cache
+                    )
+
+                action = unwrapped_model.action_readout.sample(logits)
+                past_action_id = action
+                action = action.squeeze()
+
+                # GRPO collection
+
+                meta_output = cache.prev_hiddens.meta_controller
+                old_log_probs = unwrapped_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)
+
+                next_state, reward, terminated, truncated, *_ = env.step(action.cpu().numpy())
+
+                total_reward += reward
+                step_rewards.append(reward)
+                done = terminated or truncated
+
+                if done:
+                    episode_len = step + 1
+                    break
+
+                state = next_state
+
+            # store episode
+
+            all_episodes.append((
+                cat(states, dim = 1).squeeze(0),
+                cat(log_probs, dim = 1).squeeze(0),
+                cat(switch_betas, dim = 1).squeeze(0),
+                cat(latent_actions, dim = 1).squeeze(0)
+            ))
+
+            all_cumulative_rewards.append(tensor(total_reward))
+            all_step_rewards.append(tensor(step_rewards))
+            all_episode_lens.append(episode_len)
+
+        # compute advantages
+
+        cumulative_rewards = stack(all_cumulative_rewards)
+        episode_lens = tensor(all_episode_lens)
+
+        # pad step rewards
+
+        max_len = max(all_episode_lens)
+        padded_step_rewards = torch.zeros(num_episodes, max_len)
+
+        for i, (rewards, length) in enumerate(zip(all_step_rewards, all_episode_lens)):
+            padded_step_rewards[i, :length] = rewards
+
+        # reward shaping hook
+
+        shaped_rewards = reward_shaping_fn(cumulative_rewards, padded_step_rewards, episode_lens)
+
+        if not exists(shaped_rewards):
+            continue
+
+        group_advantages = z_score(shaped_rewards)
+
+        group_states, group_log_probs, group_switch_betas, group_latent_actions = zip(*all_episodes)
+
+        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(
+                states = states,
+                log_probs = log_probs,
+                switch_betas = switch_betas,
+                latent_actions = latent_actions,
+                advantages = advantages
+            )
+
+        # learn
+
+        if len(replay_buffer) >= buffer_size:
+            dl = replay_buffer.dataloader(batch_size = num_groups)
+            dl = accelerator.prepare(dl)
+
+            meta_controller.train()
+
+            batch = next(iter(dl))
+
+            loss = meta_controller.policy_loss(
+                batch['states'],
+                batch['log_probs'],
+                batch['latent_actions'],
+                batch['advantages'],
+                batch['switch_betas'] == 1.,
+                episode_lens = batch['_lens']
+            )
+
+            accelerator.backward(loss)
+
+            grad_norm = accelerator.clip_grad_norm_(meta_controller.parameters(), max_grad_norm)
+
+            optim.step()
+            optim.zero_grad()
+
+            meta_controller.eval()
+
+            pbar.set_postfix(
+                loss = f'{loss.item():.4f}',
+                grad_norm = f'{grad_norm.item():.4f}',
+                reward = f'{cumulative_rewards.mean().item():.4f}'
+            )
+
+            accelerator.log({
+                'loss': loss.item(),
+                'grad_norm': grad_norm.item()
+            })
+
+            accelerator.print(f'loss: {loss.item():.4f}, grad_norm: {grad_norm.item():.4f}')
+
+    env.close()
+
+    # save
+
+    if exists(output_meta_controller_path):
+        unwrapped_meta_controller.save(output_meta_controller_path)
+        accelerator.print(f'MetaController weights saved to {output_meta_controller_path}')
+
+if __name__ == '__main__':
+    Fire(main)

{metacontroller_pytorch-0.0.37 → metacontroller_pytorch-0.0.40}/train_behavior_clone_babyai.py

@@ -92,8 +92,8 @@ def train(
     else: state_dim = int(torch.tensor(state_shape).prod().item())
 
     # deduce num_actions from the environment
-    minigrid.register_minigrid_envs()
-    temp_env = gym.make(env_id)
+    from babyai_env import create_env
+    temp_env = create_env(env_id)
     num_actions = int(temp_env.action_space.n)
     temp_env.close()
 

metacontroller_pytorch-0.0.37/test_babyai_e2e.sh

@@ -1,14 +0,0 @@
-#!/bin/bash
-set -e
-
-# 1. Gather trajectories
-echo "Gathering trajectories..."
-uv run gather_babyai_trajs.py --num_seeds 1000 --num_episodes_per_seed 100 --output_dir end_to_end_trajectories --env_id BabyAI-MiniBossLevel-v0
-
-# 2. Behavioral cloning
-echo "Training behavioral cloning model..."
-uv run train_behavior_clone_babyai.py --cloning_epochs 10 --discovery_epochs 10 --batch_size 256 --input_dir end_to_end_trajectories --env_id BabyAI-MiniBossLevel-v0 --checkpoint_path end_to_end_model.pt --use_resnet
-
-# 3. Inference rollouts
-echo "Running inference rollouts..."
-uv run train_babyai.py --weights_path end_to_end_model.pt --env_name BabyAI-MiniBossLevel-v0 --num_episodes 5 --buffer_size 100 --max_timesteps 100

metacontroller_pytorch-0.0.37/train_babyai.py

@@ -1,140 +0,0 @@
-# /// script
-# dependencies = [
-#   "fire",
-#   "gymnasium",
-#   "gymnasium[other]",
-#   "memmap-replay-buffer>=0.0.12",
-#   "metacontroller-pytorch",
-#   "minigrid",
-#   "tqdm"
-# ]
-# ///
-
-from fire import Fire
-from tqdm import tqdm
-from shutil import rmtree
-from pathlib import Path
-
-import torch
-from einops import rearrange
-
-import gymnasium as gym
-import minigrid
-from minigrid.wrappers import FullyObsWrapper, SymbolicObsWrapper
-
-from memmap_replay_buffer import ReplayBuffer
-from metacontroller.metacontroller import Transformer
-
-# functions
-
-def exists(v):
-    return v is not None
-
-def default(v, d):
-    return v if exists(v) else d
-
-def divisible_by(num, den):
-    return (num % den) == 0
-
-# main
-
-def main(
-    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 = None,
-    weights_path = None
-):
-
-    # environment
-
-    env = gym.make(env_name, render_mode = 'rgb_array')
-    env = FullyObsWrapper(env.unwrapped)
-    env = SymbolicObsWrapper(env.unwrapped)
-
-    rmtree(video_folder, ignore_errors = True)
-
-    env = gym.wrappers.RecordVideo(
-        env = env,
-        video_folder = video_folder,
-        name_prefix = 'babyai',
-        episode_trigger = lambda eps_num: divisible_by(eps_num, render_every_eps),
-        disable_logger = True
-    )
-
-    # maybe load model
-
-    model = None
-    if exists(weights_path):
-        weights_path = Path(weights_path)
-        assert weights_path.exists(), f"weights not found at {weights_path}"
-        model = Transformer.init_and_load(str(weights_path), strict = False)
-        model.eval()
-
-    # replay
-
-    replay_buffer = ReplayBuffer(
-        './replay-data',
-        max_episodes = buffer_size,
-        max_timesteps = max_timesteps + 1,
-        fields = dict(
-            action = 'int',
-            state_image = ('float', (7, 7, 3)),
-            state_direction = 'int'
-        ),
-        overwrite = True,
-        circular = True
-    )
-
-    # rollouts
-
-    for _ in tqdm(range(num_episodes)):
-
-        state, *_ = env.reset(seed = seed)
-
-        cache = None
-        past_action_id = None
-
-        for _ in range(max_timesteps):
-
-            if exists(model):
-                # preprocess state
-                # assume state is a dict with 'image'
-                image = state['image']
-                image_tensor = torch.from_numpy(image).float()
-                image_tensor = rearrange(image_tensor, 'h w c -> 1 1 (h w c)')
-
-                if exists(past_action_id) and 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 = model.action_readout.sample(logits)
-                past_action_id = action
-                action = action.squeeze()
-            else:
-                action = torch.randint(0, 7, ())
-
-            next_state, reward, terminated, truncated, *_ = env.step(action.cpu().numpy())
-
-            done = terminated or truncated
-
-            if done:
-                break
-
-            state = next_state
-
-    env.close()
-
-if __name__ == '__main__':
-    Fire(main)