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

{metacontroller_pytorch-0.0.43 → metacontroller_pytorch-0.0.46}/.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.46}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.43
+Version: 0.0.46
 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.46}/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.46}/metacontroller/metacontroller.py

@@ -54,6 +54,19 @@ def default(*args):
 def straight_through(src, tgt):
     return tgt + src - src.detach()
 
+# losses
+
+BehavioralCloningLosses = namedtuple('BehavioralCloningLosses', (
+    'state',
+    'action'
+))
+
+DiscoveryLosses = namedtuple('DiscoveryLosses', (
+    'action_recon',
+    'kl',
+    'switch'
+))
+
 # meta controller
 
 MetaControllerOutput = namedtuple('MetaControllerOutput', (
@@ -407,10 +420,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,
@@ -441,12 +463,15 @@ class Transformer(Module):
         return_raw_action_dist = False,
         return_latents = False,
         return_cache = False,
-        episode_lens: Tensor | None = None
+        episode_lens: Tensor | None = None,
+        return_meta_controller_output = False
     ):
         device = state.device
 
         # 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:
@@ -533,13 +558,23 @@ class Transformer(Module):
 
             action_clone_loss = self.action_readout.calculate_loss(dist_params, target_actions, mask = loss_mask)
 
-            return state_clone_loss, action_clone_loss
+            losses = BehavioralCloningLosses(state_clone_loss, action_clone_loss)
+
+            if not return_meta_controller_output:
+                return losses
+
+            return losses, next_meta_hiddens
 
         elif discovery_phase:
 
             action_recon_loss = self.action_readout.calculate_loss(dist_params, target_actions)
 
-            return action_recon_loss, next_meta_hiddens.kl_loss, next_meta_hiddens.switch_loss
+            losses = DiscoveryLosses(action_recon_loss, next_meta_hiddens.kl_loss, next_meta_hiddens.switch_loss)
+
+            if not return_meta_controller_output:
+                return losses
+
+            return losses, next_meta_hiddens
 
         # returning
 

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

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

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

@@ -6,7 +6,8 @@
 #   "memmap-replay-buffer>=0.0.12",
 #   "metacontroller-pytorch",
 #   "minigrid",
-#   "tqdm"
+#   "tqdm",
+#   "wandb"
 # ]
 # ///
 
@@ -17,35 +18,46 @@ from shutil import rmtree
 from tqdm import tqdm
 
 import torch
-from torch import cat, tensor, stack
+from torch import cat, tensor, stack, Tensor
 from torch.optim import Adam
 
 from einops import rearrange
 
+from torch_einops_utils import pad_sequence
+
 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, extract_grpo_data
 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:
+    cumulative_rewards: Tensor, # float(num_episodes,)
+    all_rewards: Tensor,        # float(num_episodes, max_timesteps)
+    episode_lens: Tensor,       # int(num_episodes,)
+    reject_threshold_cumulative_reward_variance: float = 0.
+) -> 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,)
     """
+
+    if cumulative_rewards.var() < reject_threshold_cumulative_reward_variance:
+        return None
+
     return cumulative_rewards
 
+def should_reject_group_based_on_switch_betas(
+    switch_betas: Tensor,
+    episode_lens: Tensor
+):
+    return switch_betas.sum().item() == 0.
+
 # helpers
 
 def exists(v):
@@ -68,12 +80,14 @@ def main(
     meta_controller_weights_path: str | None = None,
     output_meta_controller_path = 'metacontroller_rl_trained.pt',
     use_resnet = False,
+    num_epochs = 3,
     lr = 1e-4,
     batch_size = 16,
     num_groups = 16,
     max_grad_norm = 1.0,
     use_wandb = False,
-    wandb_project = 'metacontroller-babyai-rl'
+    wandb_project = 'metacontroller-babyai-rl',
+    reject_threshold_cumulative_reward_variance = 0.
 ):
     # accelerator
 
@@ -233,25 +247,37 @@ def main(
         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
+        # pad step rewards
+
+        padded_step_rewards = pad_sequence(all_step_rewards, dim = 0)
 
         # reward shaping hook
 
-        shaped_rewards = reward_shaping_fn(cumulative_rewards, padded_step_rewards, episode_lens)
+        shaped_rewards = reward_shaping_fn(
+            cumulative_rewards,
+            padded_step_rewards,
+            episode_lens,
+            reject_threshold_cumulative_reward_variance = reject_threshold_cumulative_reward_variance
+        )
 
         if not exists(shaped_rewards):
+            accelerator.print(f'group rejected - variance of {cumulative_rewards.var().item():.4f} is lower than threshold of {reject_threshold_cumulative_reward_variance}')
             continue
 
         group_advantages = z_score(shaped_rewards)
 
         group_states, group_log_probs, group_switch_betas, group_latent_actions = zip(*all_episodes)
 
+        # whether to reject group based on switch betas (as it determines the mask for learning)
+
+        padded_group_switch_betas, episode_lens = pad_sequence(group_switch_betas, dim = 0, return_lens = True)
+
+        if should_reject_group_based_on_switch_betas(padded_group_switch_betas, episode_lens):
+            accelerator.print(f'group rejected - switch betas for the entire group does not meet criteria for learning')
+            continue
+
         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,
@@ -264,43 +290,45 @@ def main(
         # learn
 
         if len(replay_buffer) >= buffer_size:
-            dl = replay_buffer.dataloader(batch_size = batch_size)
+            dl = replay_buffer.dataloader(batch_size = batch_size, shuffle = True)
             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']
-            )
+            for epoch in range(num_epochs):
+                for batch in 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)
+                    accelerator.backward(loss)
 
-            grad_norm = accelerator.clip_grad_norm_(meta_controller.parameters(), max_grad_norm)
+                    grad_norm = accelerator.clip_grad_norm_(meta_controller.parameters(), max_grad_norm)
 
-            optim.step()
-            optim.zero_grad()
+                    optim.step()
+                    optim.zero_grad()
 
-            meta_controller.eval()
+                    pbar.set_postfix(
+                        epoch = epoch,
+                        loss = f'{loss.item():.4f}',
+                        grad_norm = f'{grad_norm.item():.4f}',
+                        reward = f'{cumulative_rewards.mean().item():.4f}'
+                    )
 
-            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(),
+                        'reward': cumulative_rewards.mean().item()
+                    })
 
-            accelerator.log({
-                'loss': loss.item(),
-                'grad_norm': grad_norm.item()
-            })
+            meta_controller.eval()
 
-            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.46}/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)
@@ -87,11 +110,13 @@ def train(
 
     # state shape and action dimension
     # state: (B, T, H, W, C) or (B, T, D)
+
     state_shape = replay_buffer.shapes['state']
     if use_resnet: state_dim = 256
     else: state_dim = int(torch.tensor(state_shape).prod().item())
 
     # deduce num_actions from the environment
+
     from babyai_env import create_env
     temp_env = create_env(env_id)
     num_actions = int(temp_env.action_space.n)
@@ -99,6 +124,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 +137,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 +153,7 @@ def train(
 
     # training
 
+    gradient_step = 0
     for epoch in range(cloning_epochs + discovery_epochs):
 
         model.train()
@@ -154,13 +181,15 @@ 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(
+                losses, meta_controller_output = 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,
+                    return_meta_controller_output = True
                 )
 
                 if is_discovering:
@@ -175,7 +204,8 @@ def train(
                     log = dict(
                         action_recon_loss = action_recon_loss.item(),
                         kl_loss = kl_loss.item(),
-                        switch_loss = switch_loss.item()
+                        switch_loss = switch_loss.item(),
+                        switch_density = meta_controller_output.switch_beta.mean().item()
                     )
                 else:
                     state_loss, action_loss = losses
@@ -190,14 +220,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 +246,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