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

{metacontroller_pytorch-0.0.43 → metacontroller_pytorch-0.0.44}/.gitignore

@@ -1,5 +1,11 @@
 replay-data/
 recordings/
+trajectories/
+wandb/
+checkpoints/
+*.sh
+*.out
+*.slurm
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

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

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

{metacontroller_pytorch-0.0.43 → metacontroller_pytorch-0.0.44}/gather_babyai_trajs.py

@@ -37,6 +37,9 @@ from minigrid.core.constants import OBJECT_TO_IDX
 
 from memmap_replay_buffer import ReplayBuffer
 
+# Difficulty thresholds based on mission length
+EASY_MAX_LENGTH = 30      # easy: 0 to 30
+MEDIUM_MAX_LENGTH = 75    # medium: 30 to 75, hard: > 75
 
 # helpers
 
@@ -46,6 +49,67 @@ def exists(val):
 def sample(prob):
     return random.random() < prob
 
+def get_mission_length(env_id, seed):
+    """
+    Get the mission length for a given seed.
+    Returns the length of the mission string.
+    """
+    env = gym.make(env_id, render_mode="rgb_array")
+    env.reset(seed=seed)
+    length = len(env.unwrapped.mission)
+    env.close()
+    return length
+
+def categorize_seeds_by_difficulty(env_id, num_seeds_per_level, level_difficulty=None):
+    """
+    Scan seeds and categorize them by difficulty based on mission length.
+    
+    Args:
+        env_id: Environment ID
+        num_seeds_per_level: Number of seeds needed per difficulty level
+        level_difficulty: List of levels to collect seeds for.
+                       Supported: 'easy', 'medium', 'hard'
+                       If None, collects for ['easy', 'hard'].
+        max_seed_to_scan: Maximum seed value to scan
+    
+    Returns:
+        dict with keys for each requested level, each containing a list of seeds
+    """
+    
+    seeds = {level: [] for level in level_difficulty}
+    
+    total_needed = sum(num_seeds_per_level for _ in level_difficulty)
+    print(f"Scanning seeds to categorize by difficulty (need {num_seeds_per_level} per level for {level_difficulty})...")
+    
+    with tqdm(total=total_needed, desc="Categorizing seeds") as pbar:
+        seed = 1
+        all_done = False
+        while not all_done:
+            # Check if we have enough seeds for all requested levels
+            all_done = all(len(seeds[level]) >= num_seeds_per_level for level in level_difficulty)
+            
+            try:
+                mission_length = get_mission_length(env_id, seed)
+                
+                # easy: mission length <= 30
+                if 'easy' in level_difficulty and mission_length <= EASY_MAX_LENGTH and len(seeds['easy']) < num_seeds_per_level:
+                    seeds['easy'].append(seed)
+                    pbar.update(1)
+                # medium: mission length <= 75 (combines easy and medium)
+                elif 'medium' in level_difficulty and mission_length <= MEDIUM_MAX_LENGTH and len(seeds['medium']) < num_seeds_per_level:
+                    seeds['medium'].append(seed)
+                    pbar.update(1)
+                # hard: mission length > 75
+                elif 'hard' in level_difficulty and mission_length > MEDIUM_MAX_LENGTH and len(seeds['hard']) < num_seeds_per_level:
+                    seeds['hard'].append(seed)
+                    pbar.update(1)
+            except Exception as e:
+                logger.warning(f"Error getting mission length for seed {seed}: {e}")
+            
+            seed += 1
+    
+    return seeds
+
 # wrapper, necessarily modified to allow for both rgb obs (policy) and symbolic obs (bot)
 
 class RGBImgPartialObsWrapper(ObservationWrapper):
@@ -128,7 +192,7 @@ def collect_single_episode(env_id, seed, num_steps, random_action_prob, state_sh
                 env.close()
                 return None, None, False, 0
 
-            episode_state[_step] = state_obs["rgb_image"] / 255. # normalizd to 0 to 1
+            episode_state[_step] = state_obs["rgb_image"]
             episode_action[_step] = action
 
             state_obs, reward, terminated, truncated, info = env.step(action)
@@ -151,6 +215,7 @@ def collect_demonstrations(
     num_steps = 500,
     random_action_prob = 0.05,
     num_workers = None,
+    difficulty = "easy",
     output_dir = "babyai-minibosslevel-trajectories"
 ):
     """
@@ -178,11 +243,9 @@ def collect_demonstrations(
 
     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)
+    # Collect seeds by difficulty
+    assert difficulty in ['easy', 'medium', 'hard']
+    seeds = categorize_seeds_by_difficulty(env_id, num_seeds_per_level=num_seeds, level_difficulty=[difficulty])
 
     successful = 0
     progressbar = tqdm(total=total_episodes)
@@ -203,14 +266,17 @@ def collect_demonstrations(
     )
 
     # Parallel execution with bounded pending futures to avoid OOM
-    max_pending = num_workers * 4
+    max_pending = num_workers
+
+    # Flatten seeds: repeat each seed num_episodes_per_seed times
+    all_seeds = seeds[difficulty] * num_episodes_per_seed
 
     with ProcessPoolExecutor(max_workers=num_workers) as executor:
-        seed_iter = iter(seeds)
+        seed_iter = iter(all_seeds)
         futures = {}
 
         # Initial batch of submissions
-        for _ in range(min(max_pending, len(seeds))):
+        for _ in range(min(max_pending, len(all_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)
@@ -244,7 +310,13 @@ def collect_demonstrations(
     buffer.flush()
     progressbar.close()
 
+    # Save the seeds used for reproducibility
+    seeds_array = np.array(seeds[difficulty])
+    seeds_path = output_folder / "seeds.npy"
+    np.save(seeds_path, seeds_array)
+
     logger.info(f"Saved {successful} trajectories to {output_dir}")
+    logger.info(f"Saved {len(seeds_array)} seeds to {seeds_path}")
 
 if __name__ == "__main__":
     fire.Fire(collect_demonstrations)

{metacontroller_pytorch-0.0.43 → metacontroller_pytorch-0.0.44}/metacontroller/metacontroller.py

@@ -407,10 +407,19 @@ class Transformer(Module):
 
         # meta controller
 
-        self.meta_controller = meta_controller
+        self.meta_controller = meta_controller 
 
         self.register_buffer('zero', tensor(0.), persistent = False)
 
+        # ensure devices match
+        
+        if exists(self.meta_controller): self._ensure_consistent_device(self.meta_controller)
+
+    def _ensure_consistent_device(self, network):
+        self.model_device = next(self.parameters()).device
+        if next(network.parameters()).device != self.model_device:
+            network.to(self.model_device)
+
     def evolve(
         self,
         num_generations,
@@ -447,6 +456,8 @@ class Transformer(Module):
 
         # meta controller is either given or already given at init
 
+        if exists(meta_controller): self._ensure_consistent_device(meta_controller)
+
         meta_controller = default(meta_controller, self.meta_controller)
 
         if force_behavior_cloning:

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

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

{metacontroller_pytorch-0.0.43 → metacontroller_pytorch-0.0.44}/train_babyai.py

@@ -6,7 +6,8 @@
 #   "memmap-replay-buffer>=0.0.12",
 #   "metacontroller-pytorch",
 #   "minigrid",
-#   "tqdm"
+#   "tqdm",
+#   "wandb"
 # ]
 # ///
 
@@ -297,10 +298,11 @@ def main(
 
             accelerator.log({
                 'loss': loss.item(),
-                'grad_norm': grad_norm.item()
+                'grad_norm': grad_norm.item(),
+                'reward': cumulative_rewards.mean().item()
             })
 
-            accelerator.print(f'loss: {loss.item():.4f}, grad_norm: {grad_norm.item():.4f}')
+            accelerator.print(f'loss: {loss.item():.4f}, grad_norm: {grad_norm.item():.4f}, reward: {cumulative_rewards.mean().item():.4f}')
 
     env.close()
 

{metacontroller_pytorch-0.0.43 → metacontroller_pytorch-0.0.44}/train_behavior_clone_babyai.py

@@ -18,7 +18,7 @@ from tqdm import tqdm
 from pathlib import Path
 
 import torch
-from torch.optim import Adam
+from torch.optim import AdamW
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator
@@ -31,14 +31,20 @@ from metacontroller.transformer_with_resnet import TransformerWithResnet
 import minigrid
 import gymnasium as gym
 
+# TODO: loss is still ~300 and it could be the resnet output?
+# TODO: changelog (paper hparams, checkpointing, difficulty levels in trajectory collection)
+
 def train(
     input_dir = "babyai-minibosslevel-trajectories",
     env_id = "BabyAI-MiniBossLevel-v0",
     cloning_epochs = 10,
     discovery_epochs = 10,
-    batch_size = 32,
+    batch_size = 128,
+    gradient_accumulation_steps = None, 
     lr = 1e-4,
     discovery_lr = 1e-4,
+    weight_decay = 0.03,
+    discovery_weight_decay = 0.03,
     dim = 512,
     depth = 2,
     heads = 8,
@@ -47,6 +53,7 @@ def train(
     wandb_project = "metacontroller-babyai-bc",
     checkpoint_path = "transformer_bc.pt",
     meta_controller_checkpoint_path = "meta_controller_discovery.pt",
+    save_steps = 50,
     state_loss_weight = 1.,
     action_loss_weight = 1.,
     discovery_action_recon_loss_weight = 1.,
@@ -55,6 +62,22 @@ def train(
     max_grad_norm = 1.,
     use_resnet = False
 ):
+
+    def store_checkpoint(step:int):
+        if accelerator.is_main_process:
+
+            # Add step to checkpoint filenames
+            checkpoint_path_with_step = checkpoint_path.replace('.pt', f'_step_{step}.pt')
+            meta_controller_checkpoint_path_with_step = meta_controller_checkpoint_path.replace('.pt', f'_step_{step}.pt')
+
+            unwrapped_model = accelerator.unwrap_model(model)
+            unwrapped_model.save(checkpoint_path_with_step)
+
+            unwrapped_meta_controller = accelerator.unwrap_model(meta_controller)
+            unwrapped_meta_controller.save(meta_controller_checkpoint_path_with_step)
+
+            accelerator.print(f"Model saved to {checkpoint_path_with_step}, MetaController to {meta_controller_checkpoint_path_with_step}")
+
     # accelerator
 
     accelerator = Accelerator(log_with = "wandb" if use_wandb else None)
@@ -99,6 +122,10 @@ def train(
 
     accelerator.print(f"Detected state_dim: {state_dim}, num_actions: {num_actions} from env: {env_id}")
 
+    # meta controller
+
+    meta_controller = MetaController(dim)
+
     # transformer
     
     transformer_class = TransformerWithResnet if use_resnet else Transformer
@@ -108,18 +135,15 @@ def train(
         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)
+        upper_body = dict(depth = depth, heads = heads, attn_dim_head = dim_head),
+        meta_controller = meta_controller
     )
 
-    # meta controller
-
-    meta_controller = MetaController(dim)
-
     # optimizer
 
-    optim_model = Adam(model.parameters(), lr = lr)
+    optim_model = AdamW(model.parameters(), lr = lr, weight_decay = weight_decay)
 
-    optim_meta_controller = Adam(meta_controller.discovery_parameters(), lr = discovery_lr)
+    optim_meta_controller = AdamW(meta_controller.discovery_parameters(), lr = discovery_lr, weight_decay = discovery_weight_decay)
 
     # prepare
 
@@ -127,6 +151,7 @@ def train(
 
     # training
 
+    gradient_step = 0
     for epoch in range(cloning_epochs + discovery_epochs):
 
         model.train()
@@ -154,13 +179,14 @@ def train(
             else: # flatten state: (B, T, 7, 7, 3) -> (B, T, 147)
                 states = rearrange(states, 'b t ... -> b t (...)')
 
+
             with accelerator.accumulate(model):
                 losses = model(
                     states,
                     actions,
                     episode_lens = episode_lens,
                     discovery_phase = is_discovering,
-                    meta_controller = meta_controller if is_discovering else None
+                    force_behavior_cloning = not is_discovering
                 )
 
                 if is_discovering:
@@ -190,14 +216,19 @@ def train(
                         action_loss = action_loss.item(),
                     )
 
+                # gradient accumulation
+
+                if gradient_accumulation_steps is not None: loss /= gradient_accumulation_steps
+
                 # backprop
 
                 accelerator.backward(loss)
 
                 grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm = max_grad_norm)
 
-                optim.step()
-                optim.zero_grad()
+                if gradient_accumulation_steps is None or gradient_step % gradient_accumulation_steps == 0:
+                    optim.step()
+                    optim.zero_grad()
 
             # log
             
@@ -211,23 +242,21 @@ def train(
             })
 
             progress_bar.set_postfix(**log)
+            gradient_step += 1
+
+            # checkpoint 
+
+            if gradient_step % save_steps == 0:
+                accelerator.wait_for_everyone()
+                store_checkpoint(gradient_step)
 
         avg_losses = {k: v / len(dataloader) for k, v in total_losses.items()}
         avg_losses_str = ", ".join([f"{k}={v:.4f}" for k, v in avg_losses.items()])
         accelerator.print(f"Epoch {epoch}: {avg_losses_str}")
 
     # save weights
-
     accelerator.wait_for_everyone()
-    if accelerator.is_main_process:
-
-        unwrapped_model = accelerator.unwrap_model(model)
-        unwrapped_model.save(checkpoint_path)
-
-        unwrapped_meta_controller = accelerator.unwrap_model(meta_controller)
-        unwrapped_meta_controller.save(meta_controller_checkpoint_path)
-
-        accelerator.print(f"Model saved to {checkpoint_path}, MetaController to {meta_controller_checkpoint_path}")
+    store_checkpoint(gradient_step)
 
     accelerator.end_training()
 

metacontroller_pytorch-0.0.43/test_babyai_e2e.sh

@@ -1,35 +0,0 @@
-#!/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