metacontroller-pytorch 0.0.24__tar.gz → 0.0.26__tar.gz

{metacontroller_pytorch-0.0.24 → metacontroller_pytorch-0.0.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.24
+Version: 0.0.26
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller

metacontroller_pytorch-0.0.26/gather_babyai_trajs.py

@@ -0,0 +1,209 @@
+# /// script
+# dependencies = [
+#   "gymnasium",
+#   "minigrid",
+#   "tqdm",
+#   "fire",
+#   "memmap-replay-buffer>=0.0.23",
+#   "loguru"
+# ]
+# ///
+
+# taken with modifications from https://github.com/ddidacus/bot-minigrid-babyai/blob/main/tests/get_trajectories.py
+
+import fire
+import random
+import multiprocessing
+from loguru import logger
+from pathlib import Path
+
+import warnings
+warnings.filterwarnings("ignore", category = UserWarning)
+
+import numpy as np
+
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor, wait, FIRST_COMPLETED
+
+import torch
+import minigrid
+import gymnasium as gym
+from minigrid.utils.baby_ai_bot import BabyAIBot
+from minigrid.wrappers import FullyObsWrapper, SymbolicObsWrapper
+
+from memmap_replay_buffer import ReplayBuffer
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def sample(prob):
+    return random.random() < prob
+
+# agent
+
+class BabyAIBotEpsilonGreedy:
+    def __init__(self, env, random_action_prob = 0.):
+        self.expert = BabyAIBot(env)
+        self.random_action_prob = random_action_prob
+        self.num_actions = env.action_space.n
+        self.last_action = None
+
+    def __call__(self, state):
+        if sample(self.random_action_prob):
+            action = torch.randint(0, self.num_actions, ()).item()
+        else:
+            action = self.expert.replan(self.last_action)
+
+        self.last_action = action
+        return action
+
+# functions
+
+def collect_single_episode(env_id, seed, num_steps, random_action_prob, state_shape):
+    """
+    Collect a single episode of demonstrations.
+    Returns tuple of (episode_state, episode_action, success, episode_length)
+    """
+    if env_id not in gym.envs.registry:
+        minigrid.register_minigrid_envs()
+
+    env = gym.make(env_id, render_mode="rgb_array", highlight=False)
+    env = FullyObsWrapper(env.unwrapped)
+    env = SymbolicObsWrapper(env.unwrapped)
+
+    try:
+        state_obs, _ = env.reset(seed=seed)
+        
+        episode_state = np.zeros((num_steps, *state_shape), dtype=np.float32)
+        episode_action = np.zeros(num_steps, dtype=np.float32)
+
+        expert = BabyAIBotEpsilonGreedy(env.unwrapped, random_action_prob = random_action_prob)
+
+        for _step in range(num_steps):
+            try:
+                action = expert(state_obs)
+            except Exception:
+                env.close()
+                return None, None, False, 0
+
+            episode_state[_step] = state_obs["image"]
+            episode_action[_step] = action
+
+            state_obs, reward, terminated, truncated, info = env.step(action)
+
+            if terminated:
+                env.close()
+                return episode_state, episode_action, True, _step + 1
+
+        env.close()
+        return episode_state, episode_action, False, num_steps
+
+    except Exception:
+        env.close()
+        return None, None, False, 0
+
+def collect_demonstrations(
+    env_id = "BabyAI-MiniBossLevel-v0",
+    num_seeds = 100,
+    num_episodes_per_seed = 100,
+    num_steps = 500,
+    random_action_prob = 0.05,
+    num_workers = None,
+    output_dir = "babyai-minibosslevel-trajectories"
+):
+    """
+    The BabyAI Bot should be able to solve all BabyAI environments,
+    allowing us therefore to generate demonstrations.
+    Parallelized version using ProcessPoolExecutor.
+    """
+
+    # Register minigrid envs if not already registered
+    if env_id not in gym.envs.registry:
+        minigrid.register_minigrid_envs()
+
+    # Determine state shape from environment
+    temp_env = gym.make(env_id)
+    temp_env = FullyObsWrapper(temp_env.unwrapped)
+    temp_env = SymbolicObsWrapper(temp_env.unwrapped)
+    state_shape = temp_env.observation_space['image'].shape
+    temp_env.close()
+
+    logger.info(f"Detected state shape: {state_shape} for env {env_id}")
+
+    if not exists(num_workers):
+        num_workers = multiprocessing.cpu_count()
+
+    total_episodes = num_seeds * num_episodes_per_seed
+
+    # Prepare seeds for all episodes
+    seeds = []
+    for count in range(num_seeds):
+        for it in range(num_episodes_per_seed):
+            seeds.append(count + 1)
+
+    successful = 0
+    progressbar = tqdm(total=total_episodes)
+
+    output_folder = Path(output_dir)
+
+    fields = {
+        'state': ('float', state_shape),
+        'action': ('float', ())
+    }
+
+    buffer = ReplayBuffer(
+        folder = output_folder,
+        max_episodes = total_episodes,
+        max_timesteps = num_steps,
+        fields = fields,
+        overwrite = True
+    )
+
+    # Parallel execution with bounded pending futures to avoid OOM
+    max_pending = num_workers * 4
+
+    with ProcessPoolExecutor(max_workers=num_workers) as executor:
+        seed_iter = iter(seeds)
+        futures = {}
+
+        # Initial batch of submissions
+        for _ in range(min(max_pending, len(seeds))):
+            seed = next(seed_iter, None)
+            if exists(seed):
+                future = executor.submit(collect_single_episode, env_id, seed, num_steps, random_action_prob, state_shape)
+                futures[future] = seed
+
+        # Process completed tasks and submit new ones
+        while futures:
+            # Wait for at least one future to complete
+            done, _ = wait(futures, return_when=FIRST_COMPLETED)
+
+            for future in done:
+                del futures[future]
+                episode_state, episode_action, success, episode_length = future.result()
+
+                if success and exists(episode_state):
+                    buffer.store_episode(
+                        state = episode_state[:episode_length],
+                        action = episode_action[:episode_length]
+                    )
+                    successful += 1
+
+                progressbar.update(1)
+                progressbar.set_description(f"success rate = {successful}/{progressbar.n:.2f}")
+
+                # Submit a new task to replace the completed one
+                seed = next(seed_iter, None)
+                if exists(seed):
+                    new_future = executor.submit(collect_single_episode, env_id, seed, num_steps, random_action_prob, state_shape)
+                    futures[new_future] = seed
+
+    buffer.flush()
+    progressbar.close()
+
+    logger.info(f"Saved {successful} trajectories to {output_dir}")
+
+if __name__ == "__main__":
+    fire.Fire(collect_demonstrations)

{metacontroller_pytorch-0.0.24 → metacontroller_pytorch-0.0.26}/metacontroller/metacontroller.py

@@ -6,7 +6,7 @@ from collections import namedtuple
 from loguru import logger
 
 import torch
-from torch import nn, cat, stack, tensor
+from torch import nn, cat, stack, tensor, Tensor
 from torch.nn import Module, GRU, Linear, Identity
 import torch.nn.functional as F
 
@@ -26,7 +26,7 @@ from discrete_continuous_embed_readout import Embed, Readout, EmbedAndReadout
 
 from assoc_scan import AssocScan
 
-from torch_einops_utils import pad_at_dim
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask
 from torch_einops_utils.save_load import save_load
 
 # constants
@@ -151,7 +151,8 @@ class MetaController(Module):
         cache: MetaControllerOutput | None = None,
         discovery_phase = False,
         hard_switch = False,
-        temperature = 1.
+        temperature = 1.,
+        episode_lens: Tensor | None = None
     ):
         device = residual_stream.device
 
@@ -168,7 +169,9 @@ class MetaController(Module):
         if discovery_phase:
             logger.warning('meta controller cache being passed back in for discovery phase, which does not make sense given bidirectional encoder')
 
-            encoded_temporal = self.bidirectional_temporal_encoder(meta_embed)
+            mask = maybe(lens_to_mask)(episode_lens, meta_embed.shape[1])
+
+            encoded_temporal = self.bidirectional_temporal_encoder(meta_embed, mask = mask)
 
             proposed_action_hidden, _ = self.emitter(cat((encoded_temporal, meta_embed), dim = -1))
             readout = self.emitter_to_action_mean_log_var
@@ -335,6 +338,7 @@ class Transformer(Module):
         return_raw_action_dist = False,
         return_latents = False,
         return_cache = False,
+        episode_lens: Tensor | None = None
     ):
         device = state.device
 
@@ -362,6 +366,9 @@ class Transformer(Module):
             state, target_state = state[:, :-1], state[:, 1:]
             actions, target_actions = actions[:, :-1], actions[:, 1:]
 
+            if exists(episode_lens):
+                episode_lens = (episode_lens - 1).clamp(min = 0)
+
         # transformer lower body
 
         with lower_transformer_context():
@@ -387,7 +394,7 @@ class Transformer(Module):
         with meta_controller_context():
 
             if exists(meta_controller):
-                control_signal, next_meta_hiddens = meta_controller(residual_stream, cache = meta_hiddens, discovery_phase = discovery_phase, temperature = meta_controller_temperature)
+                control_signal, next_meta_hiddens = meta_controller(residual_stream, cache = meta_hiddens, discovery_phase = discovery_phase, temperature = meta_controller_temperature, episode_lens = episode_lens)
             else:
                 control_signal, next_meta_hiddens = self.zero, None
 
@@ -406,10 +413,14 @@ class Transformer(Module):
         # maybe return behavior cloning loss
 
         if behavioral_cloning:
+            loss_mask = None
+            if exists(episode_lens):
+                loss_mask = lens_to_mask(episode_lens, state.shape[1])
+
             state_dist_params = self.state_readout(attended)
-            state_clone_loss = self.state_readout.calculate_loss(state_dist_params, target_state)
+            state_clone_loss = self.state_readout.calculate_loss(state_dist_params, target_state, mask = loss_mask)
 
-            action_clone_loss = self.action_readout.calculate_loss(dist_params, target_actions)
+            action_clone_loss = self.action_readout.calculate_loss(dist_params, target_actions, mask = loss_mask)
 
             return state_clone_loss, action_clone_loss
 

{metacontroller_pytorch-0.0.24 → metacontroller_pytorch-0.0.26}/pyproject.toml

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

metacontroller_pytorch-0.0.26/test_babyai_e2e.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+set -e
+
+# 1. Gather trajectories
+echo "Gathering trajectories..."
+uv run gather_babyai_trajs.py --num_seeds 10 --num_episodes_per_seed 10 --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 --epochs 1 --batch_size 16 --input_dir end_to_end_trajectories --env_id BabyAI-MiniBossLevel-v0 --checkpoint_path end_to_end_model.pt
+
+# 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.24 → metacontroller_pytorch-0.0.26}/tests/test_metacontroller.py

@@ -10,12 +10,15 @@ from einops import rearrange
 
 @param('action_discrete', (False, True))
 @param('switch_per_latent_dim', (False, True))
+@param('variable_length', (False, True))
 def test_metacontroller(
     action_discrete,
-    switch_per_latent_dim
+    switch_per_latent_dim,
+    variable_length
 ):
 
     state = torch.randn(1, 1024, 384)
+    episode_lens = torch.tensor([512]) if variable_length else None
 
     if action_discrete:
         actions = torch.randint(0, 4, (1, 1024))
@@ -36,7 +39,7 @@ def test_metacontroller(
         upper_body = dict(depth = 2,),
     )
 
-    state_clone_loss, action_clone_loss = model(state, actions)
+    state_clone_loss, action_clone_loss = model(state, actions, episode_lens = episode_lens)
     (state_clone_loss + 0.5 * action_clone_loss).backward()
 
     # discovery and internal rl phase with meta controller
@@ -50,7 +53,7 @@ def test_metacontroller(
 
     # discovery phase
 
-    (action_recon_loss, kl_loss, switch_loss) = model(state, actions, meta_controller = meta_controller, discovery_phase = True)
+    (action_recon_loss, kl_loss, switch_loss) = model(state, actions, meta_controller = meta_controller, discovery_phase = True, episode_lens = episode_lens)
     (action_recon_loss + kl_loss * 0.1 + switch_loss * 0.2).backward()
 
     # internal rl - done iteratively

metacontroller_pytorch-0.0.24/train.py → metacontroller_pytorch-0.0.26/train_babyai.py

@@ -13,13 +13,17 @@
 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
 
@@ -41,12 +45,15 @@ def main(
     buffer_size = 5_000,
     render_every_eps = 1_000,
     video_folder = './recordings',
-    seed = None
+    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)
 
@@ -58,6 +65,15 @@ def main(
         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(
@@ -79,10 +95,37 @@ def main(
 
         state, *_ = env.reset(seed = seed)
 
+        cache = None
+        past_action_id = None
+
         for _ in range(max_timesteps):
 
-            action = torch.randint(0, 7, ())
-            next_state, reward, terminated, truncated, *_ = env.step(action.numpy())
+            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
 
@@ -91,7 +134,7 @@ def main(
 
             state = next_state
 
-# running
+    env.close()
 
 if __name__ == '__main__':
     Fire(main)

metacontroller_pytorch-0.0.26/train_behavior_clone_babyai.py

@@ -0,0 +1,175 @@
+# /// script
+# dependencies = [
+#   "accelerate",
+#   "fire",
+#   "memmap-replay-buffer>=0.0.23",
+#   "metacontroller-pytorch",
+#   "torch",
+#   "einops",
+#   "tqdm",
+#   "wandb",
+#   "gymnasium",
+#   "minigrid"
+# ]
+# ///
+
+import fire
+from tqdm import tqdm
+from pathlib import Path
+
+import torch
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+
+from accelerate import Accelerator
+from memmap_replay_buffer import ReplayBuffer
+from einops import rearrange
+
+from metacontroller.metacontroller import Transformer
+
+import minigrid
+import gymnasium as gym
+
+def train(
+    input_dir: str = "babyai-minibosslevel-trajectories",
+    env_id: str = "BabyAI-MiniBossLevel-v0",
+    epochs: int = 10,
+    batch_size: int = 32,
+    lr: float = 1e-4,
+    dim: int = 512,
+    depth: int = 2,
+    heads: int = 8,
+    dim_head: int = 64,
+    use_wandb: bool = False,
+    wandb_project: str = "metacontroller-babyai-bc",
+    checkpoint_path: str = "transformer_bc.pt",
+    state_loss_weight: float = 1.,
+    action_loss_weight: float = 1.
+):
+    # accelerator
+
+    accelerator = Accelerator(log_with = "wandb" if use_wandb else None)
+
+    if use_wandb:
+        accelerator.init_trackers(
+            wandb_project,
+            config = {
+                "epochs": epochs,
+                "batch_size": batch_size,
+                "lr": lr,
+                "dim": dim,
+                "depth": depth,
+                "heads": heads,
+                "dim_head": dim_head,
+                "env_id": env_id,
+                "state_loss_weight": state_loss_weight,
+                "action_loss_weight": action_loss_weight
+            }
+        )
+
+    # replay buffer and dataloader
+
+    input_path = Path(input_dir)
+    assert input_path.exists(), f"Input directory {input_dir} does not exist"
+
+    replay_buffer = ReplayBuffer.from_folder(input_path)
+    dataloader = replay_buffer.dataloader(batch_size = batch_size)
+
+    # state shape and action dimension
+    # state: (B, T, H, W, C) or (B, T, D)
+    state_shape = replay_buffer.shapes['state']
+    state_dim = int(torch.tensor(state_shape).prod().item())
+
+    # state shape and action dimension
+    # state: (B, T, H, W, C) or (B, T, D)
+    state_shape = replay_buffer.shapes['state']
+    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)
+    num_actions = int(temp_env.action_space.n)
+    temp_env.close()
+
+    accelerator.print(f"Detected state_dim: {state_dim}, num_actions: {num_actions} from env: {env_id}")
+
+    # transformer
+
+    model = Transformer(
+        dim = dim,
+        state_embed_readout = dict(num_continuous = state_dim),
+        action_embed_readout = dict(num_discrete = num_actions),
+        lower_body = dict(depth = depth, heads = heads, attn_dim_head = dim_head),
+        upper_body = dict(depth = depth, heads = heads, attn_dim_head = dim_head)
+    )
+
+    # optimizer
+
+    optim = Adam(model.parameters(), lr = lr)
+
+    # prepare
+
+    model, optim, dataloader = accelerator.prepare(model, optim, dataloader)
+
+    # training
+
+    for epoch in range(epochs):
+        model.train()
+        total_state_loss = 0.
+        total_action_loss = 0.
+
+        progress_bar = tqdm(dataloader, desc = f"Epoch {epoch}", disable = not accelerator.is_local_main_process)
+
+        for batch in progress_bar:
+            # batch is a NamedTuple (e.g. MemoryMappedBatch)
+            # state: (B, T, 7, 7, 3), action: (B, T)
+
+            states = batch['state'].float()
+            actions = batch['action'].long()
+            episode_lens = batch.get('_lens')
+
+            # flatten state: (B, T, 7, 7, 3) -> (B, T, 147)
+
+            states = rearrange(states, 'b t ... -> b t (...)')
+
+            with accelerator.accumulate(model):
+                state_loss, action_loss = model(states, actions, episode_lens = episode_lens)
+                loss = state_loss * state_loss_weight + action_loss * action_loss_weight
+
+                accelerator.backward(loss)
+                optim.step()
+                optim.zero_grad()
+
+            # log
+
+            total_state_loss += state_loss.item()
+            total_action_loss += action_loss.item()
+
+            accelerator.log({
+                "state_loss": state_loss.item(),
+                "action_loss": action_loss.item(),
+                "total_loss": loss.item()
+            })
+
+            progress_bar.set_postfix(
+                state_loss = state_loss.item(),
+                action_loss = action_loss.item()
+            )
+
+        avg_state_loss = total_state_loss / len(dataloader)
+        avg_action_loss = total_action_loss / len(dataloader)
+
+        accelerator.print(f"Epoch {epoch}: state_loss={avg_state_loss:.4f}, action_loss={avg_action_loss:.4f}")
+
+    # save weights
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        unwrapped_model = accelerator.unwrap_model(model)
+        unwrapped_model.save(checkpoint_path)
+        accelerator.print(f"Model saved to {checkpoint_path}")
+
+    accelerator.end_training()
+
+if __name__ == "__main__":
+    fire.Fire(train)