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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.44
+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.44 → 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', (
@@ -450,7 +463,8 @@ 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
 
@@ -544,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.44 → metacontroller_pytorch-0.0.46}/pyproject.toml

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

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

@@ -18,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):
@@ -69,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
 
@@ -234,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,
@@ -265,42 +290,43 @@ 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(),
-                'reward': cumulative_rewards.mean().item()
-            })
+            meta_controller.eval()
 
             accelerator.print(f'loss: {loss.item():.4f}, grad_norm: {grad_norm.item():.4f}, reward: {cumulative_rewards.mean().item():.4f}')
 

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

@@ -110,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)
@@ -181,12 +183,13 @@ def train(
 
 
             with accelerator.accumulate(model):
-                losses = model(
+                losses, meta_controller_output = model(
                     states,
                     actions,
                     episode_lens = episode_lens,
                     discovery_phase = is_discovering,
-                    force_behavior_cloning = not is_discovering
+                    force_behavior_cloning = not is_discovering,
+                    return_meta_controller_output = True
                 )
 
                 if is_discovering:
@@ -201,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