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

{metacontroller_pytorch-0.0.42 → 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.42 → metacontroller_pytorch-0.0.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.42
+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.42 → 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.42 → metacontroller_pytorch-0.0.44}/metacontroller/metacontroller.py

@@ -291,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
@@ -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.42 → metacontroller_pytorch-0.0.44}/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.42 → metacontroller_pytorch-0.0.44}/pyproject.toml

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

{metacontroller_pytorch-0.0.42 → metacontroller_pytorch-0.0.44}/tests/test_metacontroller.py

@@ -114,13 +114,14 @@ def test_metacontroller(
         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)
 
             # extract grpo data and store
 
-            grpo_data_list.append(extract_grpo_data(meta_controller, cache))
+            grpo_data = extract_grpo_data(meta_controller, cache)
+            grpo_data_list.append(grpo_data)
 
         # accumulate across time for the episode data
 
@@ -145,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.42 → metacontroller_pytorch-0.0.44}/train_babyai.py

@@ -6,7 +6,8 @@
 #   "memmap-replay-buffer>=0.0.12",
 #   "metacontroller-pytorch",
 #   "minigrid",
-#   "tqdm"
+#   "tqdm",
+#   "wandb"
 # ]
 # ///
 
@@ -57,22 +58,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
 
@@ -263,7 +265,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()
@@ -296,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.44/train_babyai_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.42 → 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.42/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