metacontroller-pytorch 0.0.35__tar.gz → 0.0.36__tar.gz

{metacontroller_pytorch-0.0.35 → metacontroller_pytorch-0.0.36}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.35
+Version: 0.0.36
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -39,7 +39,7 @@ Requires-Dist: discrete-continuous-embed-readout>=0.1.12
 Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: loguru
-Requires-Dist: memmap-replay-buffer>=0.0.23
+Requires-Dist: memmap-replay-buffer>=0.0.25
 Requires-Dist: torch-einops-utils>=0.0.19
 Requires-Dist: torch>=2.5
 Requires-Dist: vector-quantize-pytorch>=1.27.20

{metacontroller_pytorch-0.0.35 → metacontroller_pytorch-0.0.36}/metacontroller/metacontroller.py

@@ -336,6 +336,8 @@ class MetaController(Module):
 
         return control_signal, MetaControllerOutput(next_hiddens, residual_stream, action_dist, sampled_latent_action, switch_beta, kl_loss, switch_loss)
 
+MetaController.policy_loss = policy_loss
+
 # main transformer, which is subsumed into the environment after behavioral cloning
 
 Hiddens = namedtuple('Hiddens', (
@@ -414,17 +416,21 @@ class Transformer(Module):
     ):
         device = state.device
 
+        # meta controller is either given or already given at init
+
         meta_controller = default(meta_controller, self.meta_controller)
 
-        meta_controlling = exists(meta_controller)
+        has_meta_controller = exists(meta_controller)
 
-        behavioral_cloning = not meta_controlling and not return_raw_action_dist
+        assert not (discovery_phase and not has_meta_controller), 'meta controller must be made available during discovery phase'
+
+        behavioral_cloning = not has_meta_controller and not return_raw_action_dist
 
         # by default, if meta controller is passed in, transformer is no grad
 
-        lower_transformer_context = nullcontext if not meta_controlling else torch.no_grad
-        meta_controller_context = nullcontext if meta_controlling else torch.no_grad
-        upper_transformer_context = nullcontext if (not meta_controlling or discovery_phase) else torch.no_grad
+        lower_transformer_context = nullcontext if not has_meta_controller else torch.no_grad
+        meta_controller_context = nullcontext if has_meta_controller else torch.no_grad
+        upper_transformer_context = nullcontext if (not has_meta_controller or discovery_phase) else torch.no_grad
 
         # handle cache
 
@@ -432,7 +438,8 @@ class Transformer(Module):
 
         # handle maybe behavioral cloning
 
-        if behavioral_cloning or (meta_controlling and discovery_phase):
+        if behavioral_cloning or discovery_phase: # during behavior cloning and discovery phase, the network is predicting / reconstructing the next token
+
             assert exists(actions), f'`actions` cannot be empty when doing discovery or behavioral cloning'
 
             state, target_state = state[:, :-1], state[:, 1:]
@@ -495,7 +502,7 @@ class Transformer(Module):
 
             return state_clone_loss, action_clone_loss
 
-        elif meta_controlling and discovery_phase:
+        elif discovery_phase:
 
             action_recon_loss = self.action_readout.calculate_loss(dist_params, target_actions)
 

{metacontroller_pytorch-0.0.35 → metacontroller_pytorch-0.0.36}/metacontroller/metacontroller_with_binary_mapper.py

@@ -28,7 +28,7 @@ from torch_einops_utils.save_load import save_load
 
 from vector_quantize_pytorch import BinaryMapper
 
-from metacontroller.metacontroller import MetaControllerOutput
+from metacontroller.metacontroller import MetaControllerOutput, policy_loss
 
 # constants
 
@@ -170,7 +170,7 @@ class MetaControllerWithBinaryMapper(Module):
         action_log_probs = log_probs.gather(-1, codes)
         action_log_probs = rearrange(action_log_probs, '... 1 -> ...')
 
-        return action_log_probs.sum(dim = -1)
+        return action_log_probs
 
     def forward(
         self,
@@ -302,3 +302,5 @@ class MetaControllerWithBinaryMapper(Module):
             switch_beta = rearrange(switch_beta, '... 1 -> ...')
 
         return control_signal, MetaControllerOutput(next_hiddens, residual_stream, binary_logits, sampled_codes, switch_beta, kl_loss, switch_loss)
+
+MetaControllerWithBinaryMapper.policy_loss = policy_loss

{metacontroller_pytorch-0.0.35 → metacontroller_pytorch-0.0.36}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "metacontroller-pytorch"
-version = "0.0.35"
+version = "0.0.36"
 description = "Transformer Metacontroller"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -29,7 +29,7 @@ dependencies = [
     "einx>=0.3.0",
     "einops>=0.8.1",
     "loguru",
-    "memmap-replay-buffer>=0.0.23",
+    "memmap-replay-buffer>=0.0.25",
     "torch>=2.5",
     "torch-einops-utils>=0.0.19",
     "vector-quantize-pytorch>=1.27.20",

{metacontroller_pytorch-0.0.35 → metacontroller_pytorch-0.0.36}/tests/test_metacontroller.py

@@ -1,6 +1,7 @@
 import pytest
 param = pytest.mark.parametrize
 
+from shutil import rmtree
 from pathlib import Path
 from functools import partial
 
@@ -9,6 +10,8 @@ from torch import cat
 from metacontroller.metacontroller import Transformer, MetaController, policy_loss, z_score
 from metacontroller.metacontroller_with_binary_mapper import MetaControllerWithBinaryMapper
 
+from memmap_replay_buffer import ReplayBuffer
+
 from einops import rearrange
 
 # functions
@@ -66,6 +69,12 @@ def test_metacontroller(
             dim_latent = 128,
             switch_per_latent_dim = switch_per_latent_dim
         )
+
+        field_shapes = dict(
+            log_probs = ('float', 128),
+            switch_betas = ('float', 128 if switch_per_latent_dim else 1),
+            latent_actions = ('float', 128)
+        )
     else:
         meta_controller = MetaControllerWithBinaryMapper(
             dim_model = 512,
@@ -74,6 +83,12 @@ def test_metacontroller(
             dim_code_bits = 8, # 2 ** 8 = 256 codes
         )
 
+        field_shapes = dict(
+            log_probs = ('float', 8),
+            switch_betas = ('float', 8 if switch_per_latent_dim else 1),
+            latent_actions = ('float', 256)
+        )
+
     # discovery phase
 
     (action_recon_loss, kl_loss, switch_loss) = model(state, actions, meta_controller = meta_controller, discovery_phase = True, episode_lens = episode_lens)
@@ -81,6 +96,23 @@ def test_metacontroller(
 
     # internal rl - done iteratively
 
+    # replay buffer
+
+    test_folder = './test-buffer-for-grpo'
+
+    replay_buffer = ReplayBuffer(
+        test_folder,
+        max_episodes = 3,
+        max_timesteps = 256,
+        fields = dict(
+            states = ('float', 512),
+            **field_shapes
+        ),
+        meta_fields = dict(
+            advantages = 'float'
+        )
+    )
+
     # simulate grpo
 
     all_episodes = []
@@ -129,22 +161,34 @@ def test_metacontroller(
     # calculate advantages using z-score
 
     rewards = cat(all_rewards)
-    advantages = z_score(rewards)
+    group_advantages = z_score(rewards)
 
-    assert advantages.shape == (3,)
+    assert group_advantages.shape == (3,)
 
     # 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))
 
-    loss = policy_loss(
-        meta_controller,
-        group_states,
-        group_log_probs,
-        group_latent_actions,
-        advantages,
-        group_switch_betas == 1.,
-        episode_lens = episode_lens[:1].repeat(3) if exists(episode_lens) else None
+    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,
+            log_probs = log_probs,
+            switch_betas = switch_betas,
+            latent_actions = latent_actions,
+            advantages = advantages
+        )
+
+    dl = replay_buffer.dataloader(batch_size = 3)
+
+    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']
     )
 
     loss.backward()
@@ -167,3 +211,5 @@ def test_metacontroller(
 
     Path('./meta_controller.pt').unlink()
     Path('./trained.pt').unlink()
+
+    rmtree(test_folder, ignore_errors = True)

{metacontroller_pytorch-0.0.35 → metacontroller_pytorch-0.0.36}/train_behavior_clone_babyai.py

@@ -25,29 +25,35 @@ from accelerate import Accelerator
 from memmap_replay_buffer import ReplayBuffer
 from einops import rearrange
 
-from metacontroller.metacontroller import Transformer
+from metacontroller.metacontroller import Transformer, MetaController
 from metacontroller.transformer_with_resnet import TransformerWithResnet
 
 import minigrid
 import gymnasium as gym
 
 def train(
-    input_dir: str = "babyai-minibosslevel-trajectories",
-    env_id: str = "BabyAI-MiniBossLevel-v0",
-    cloning_epochs: int = 10,
-    discovery_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.,
-    use_resnet: bool = False
+    input_dir = "babyai-minibosslevel-trajectories",
+    env_id = "BabyAI-MiniBossLevel-v0",
+    cloning_epochs = 10,
+    discovery_epochs = 10,
+    batch_size = 32,
+    lr = 1e-4,
+    discovery_lr = 1e-4,
+    dim = 512,
+    depth = 2,
+    heads = 8,
+    dim_head = 64,
+    use_wandb = False,
+    wandb_project = "metacontroller-babyai-bc",
+    checkpoint_path = "transformer_bc.pt",
+    meta_controller_checkpoint_path = "meta_controller_discovery.pt",
+    state_loss_weight = 1.,
+    action_loss_weight = 1.,
+    discovery_action_recon_loss_weight = 1.,
+    discovery_kl_loss_weight = 1.,
+    discovery_switch_loss_weight = 1.,
+    max_grad_norm = 1.,
+    use_resnet = False
 ):
     # accelerator
 
@@ -96,6 +102,7 @@ def train(
     # transformer
     
     transformer_class = TransformerWithResnet if use_resnet else Transformer
+
     model = transformer_class(
         dim = dim,
         state_embed_readout = dict(num_continuous = state_dim),
@@ -104,23 +111,34 @@ def train(
         upper_body = dict(depth = depth, heads = heads, attn_dim_head = dim_head)
     )
 
+    # meta controller
+
+    meta_controller = MetaController(dim)
+
     # optimizer
 
-    optim = Adam(model.parameters(), lr = lr)
+    optim_model = Adam(model.parameters(), lr = lr)
+
+    optim_meta_controller = Adam(meta_controller.discovery_parameters(), lr = discovery_lr)
 
     # prepare
 
-    model, optim, dataloader = accelerator.prepare(model, optim, dataloader)
+    model, optim_model, optim_meta_controller, dataloader = accelerator.prepare(model, optim_model, optim_meta_controller, dataloader)
 
     # training
+
     for epoch in range(cloning_epochs + discovery_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)
+
         is_discovering = (epoch >= cloning_epochs) # discovery phase is BC with metacontroller tuning
 
+        optim = optim_model if not is_discovering else optim_meta_controller
+
         for batch in progress_bar:
             # batch is a NamedTuple (e.g. MemoryMappedBatch)
             # state: (B, T, 7, 7, 3), action: (B, T)
@@ -130,18 +148,53 @@ def train(
             episode_lens = batch.get('_lens')
 
             # use resnet18 to embed visual observations
+
             if use_resnet: 
                 states = model.visual_encode(states)
             else: # 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, discovery_phase=is_discovering)
-                loss = state_loss * state_loss_weight + action_loss * action_loss_weight
+                losses = model(
+                    states,
+                    actions,
+                    episode_lens = episode_lens,
+                    discovery_phase = is_discovering,
+                    meta_controller = meta_controller if is_discovering else None
+                )
+
+                if is_discovering:
+                    action_recon_loss, kl_loss, switch_loss = losses
+
+                    loss = (
+                        action_recon_loss * discovery_action_recon_loss_weight +
+                        kl_loss * discovery_kl_loss_weight +
+                        switch_loss * discovery_switch_loss_weight
+                    )
+
+                    log = dict(
+                        action_recon_loss = action_recon_loss.item(),
+                        kl_loss = kl_loss.item(),
+                        switch_loss = switch_loss.item()
+                    )
+                else:
+                    state_loss, action_loss = losses
+
+                    loss = (
+                        state_loss * state_loss_weight +
+                        action_loss * action_loss_weight
+                    )
+
+                    log = dict(
+                        state_loss = state_loss.item(),
+                        action_loss = action_loss.item(),
+                    )
+
+                # backprop
 
                 accelerator.backward(loss)
 
-                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm = max_grad_norm)
 
                 optim.step()
                 optim.zero_grad()
@@ -152,8 +205,7 @@ def train(
             total_action_loss += action_loss.item()
 
             accelerator.log({
-                "state_loss": state_loss.item(),
-                "action_loss": action_loss.item(),
+                **log,
                 "total_loss": loss.item(),
                 "grad_norm": grad_norm.item()
             })
@@ -172,9 +224,14 @@ def train(
 
     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}")
+
+        unwrapped_meta_controller = accelerator.unwrap_model(meta_controller)
+        unwrapped_meta_controller.save(meta_controller_checkpoint_path)
+
+        accelerator.print(f"Model saved to {checkpoint_path}, MetaControler to {meta_controller_checkpoint_path}")
 
     accelerator.end_training()