metacontroller-pytorch 0.0.25__tar.gz → 0.0.27__tar.gz

{metacontroller_pytorch-0.0.25 → metacontroller_pytorch-0.0.27}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.25
+Version: 0.0.27
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -42,6 +42,7 @@ Requires-Dist: loguru
 Requires-Dist: memmap-replay-buffer>=0.0.23
 Requires-Dist: torch-einops-utils>=0.0.16
 Requires-Dist: torch>=2.5
+Requires-Dist: vector-quantize-pytorch>=1.27.20
 Requires-Dist: x-evolution>=0.1.23
 Requires-Dist: x-mlps-pytorch
 Requires-Dist: x-transformers

{metacontroller_pytorch-0.0.25 → metacontroller_pytorch-0.0.27}/metacontroller/metacontroller.py

@@ -6,7 +6,7 @@ from collections import namedtuple
 from loguru import logger
 
 import torch
-from torch import nn, cat, stack, tensor
+from torch import nn, cat, stack, tensor, Tensor
 from torch.nn import Module, GRU, Linear, Identity
 import torch.nn.functional as F
 
@@ -26,7 +26,7 @@ from discrete_continuous_embed_readout import Embed, Readout, EmbedAndReadout
 
 from assoc_scan import AssocScan
 
-from torch_einops_utils import pad_at_dim, lens_to_mask
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask
 from torch_einops_utils.save_load import save_load
 
 # constants
@@ -151,7 +151,8 @@ class MetaController(Module):
         cache: MetaControllerOutput | None = None,
         discovery_phase = False,
         hard_switch = False,
-        temperature = 1.
+        temperature = 1.,
+        episode_lens: Tensor | None = None
     ):
         device = residual_stream.device
 
@@ -168,7 +169,9 @@ class MetaController(Module):
         if discovery_phase:
             logger.warning('meta controller cache being passed back in for discovery phase, which does not make sense given bidirectional encoder')
 
-            encoded_temporal = self.bidirectional_temporal_encoder(meta_embed)
+            mask = maybe(lens_to_mask)(episode_lens, meta_embed.shape[1])
+
+            encoded_temporal = self.bidirectional_temporal_encoder(meta_embed, mask = mask)
 
             proposed_action_hidden, _ = self.emitter(cat((encoded_temporal, meta_embed), dim = -1))
             readout = self.emitter_to_action_mean_log_var
@@ -391,7 +394,7 @@ class Transformer(Module):
         with meta_controller_context():
 
             if exists(meta_controller):
-                control_signal, next_meta_hiddens = meta_controller(residual_stream, cache = meta_hiddens, discovery_phase = discovery_phase, temperature = meta_controller_temperature)
+                control_signal, next_meta_hiddens = meta_controller(residual_stream, cache = meta_hiddens, discovery_phase = discovery_phase, temperature = meta_controller_temperature, episode_lens = episode_lens)
             else:
                 control_signal, next_meta_hiddens = self.zero, None
 

metacontroller_pytorch-0.0.27/metacontroller/metacontroller_with_binary_mapper.py

@@ -0,0 +1,266 @@
+from __future__ import annotations
+from contextlib import nullcontext
+
+from functools import partial
+from collections import namedtuple
+from loguru import logger
+
+import torch
+from torch import nn, cat, stack, tensor, Tensor
+from torch.nn import Module, GRU, Linear, Identity
+import torch.nn.functional as F
+
+# einops
+
+import einx
+from einops import einsum, rearrange, repeat, reduce
+from einops.layers.torch import Rearrange
+
+# external modules
+
+from x_transformers import Encoder, Decoder
+from x_mlps_pytorch import Feedforwards
+
+from assoc_scan import AssocScan
+
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask
+from torch_einops_utils.save_load import save_load
+
+from vector_quantize_pytorch import BinaryMapper
+
+# constants
+
+LinearNoBias = partial(Linear, bias = False)
+
+GRU = partial(GRU, batch_first = True)
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(*args):
+    for arg in args:
+        if exists(arg):
+            return arg
+    return None
+
+def straight_through(src, tgt):
+    return tgt + src - src.detach()
+
+# meta controller
+
+MetaControllerOutput = namedtuple('MetaControllerOutput', (
+    'prev_hiddens',
+    'action_dist',
+    'codes',
+    'kl_loss',
+    'switch_loss'
+))
+
+@save_load()
+class MetaControllerWithBinaryMapper(Module):
+    def __init__(
+        self,
+        dim_model,
+        *,
+        dim_meta_controller = 256,
+        dim_code_bits = 4,
+        switch_per_code = False,
+        decoder_expansion_factor = 2.,
+        decoder_depth = 1,
+        hypernetwork_low_rank = 16,
+        assoc_scan_kwargs: dict = dict(),
+        bidirectional_temporal_encoder_kwargs: dict = dict(
+            attn_dim_head = 32, heads = 8
+        ),
+        kl_loss_threshold = 0.
+    ):
+        super().__init__()
+        dim_meta = default(dim_meta_controller, dim_model)
+
+        self.model_to_meta = Linear(dim_model, dim_meta)
+
+        self.bidirectional_temporal_encoder = Encoder(dim = dim_meta, depth = 1, **bidirectional_temporal_encoder_kwargs)
+
+        self.emitter = GRU(dim_meta * 2, dim_meta * 2)
+        self.emitter_to_binary_logits = Linear(dim_meta * 2, dim_code_bits)
+
+        self.action_proposer = GRU(dim_meta, dim_meta)
+        self.proposer_to_binary_logits = Linear(dim_meta, dim_code_bits)
+
+        # binary mapper
+        # proposed in https://arxiv.org/abs/2510.17558 as a more stable alternative to VAE by François Fleuret
+
+        self.binary_mapper = BinaryMapper(
+            bits = dim_code_bits,
+            kl_loss_threshold = kl_loss_threshold
+        )
+
+        self.dim_code_bits = dim_code_bits
+        self.num_codes = self.binary_mapper.num_codes
+
+        # switching unit
+
+        self.switch_per_code = switch_per_code
+
+        self.switching_unit = GRU(dim_meta + self.num_codes, dim_meta)
+        self.to_switching_unit_beta = nn.Linear(dim_meta, self.num_codes if switch_per_code else 1, bias = False)
+
+        self.switch_gating = AssocScan(**assoc_scan_kwargs)
+
+        # decoder
+
+        assert hypernetwork_low_rank < self.num_codes
+
+        dim_decoder_hidden = int(self.num_codes * decoder_expansion_factor)
+
+        self.decoder = Feedforwards(
+            dim_in = self.num_codes,
+            dim = dim_decoder_hidden,
+            depth = decoder_depth,
+            dim_out = 2 * hypernetwork_low_rank * dim_model
+        )
+
+        self.to_hyper_network_weights = Rearrange('... (two d r) -> two ... d r', two = 2, r = hypernetwork_low_rank)
+
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
+    def discovery_parameters(self):
+        return [
+            *self.model_to_meta.parameters(),
+            *self.bidirectional_temporal_encoder.parameters(),
+            *self.emitter.parameters(),
+            *self.emitter_to_binary_logits.parameters(),
+            *self.binary_mapper.parameters(),
+            *self.decoder.parameters(),
+            *self.switch_gating.parameters()
+        ]
+
+    def internal_rl_parameters(self):
+        return [
+            *self.action_proposer.parameters(),
+            *self.proposer_to_binary_logits.parameters()
+        ]
+
+    def forward(
+        self,
+        residual_stream,
+        cache: MetaControllerOutput | None = None,
+        discovery_phase = False,
+        hard_switch = False,
+        temperature = 1.,
+        episode_lens: Tensor | None = None
+    ):
+        device = residual_stream.device
+
+        # destruct prev cache
+
+        prev_action_proposer_hidden, prev_switching_unit_gru_hidden, prev_switch_gated_hiddens, prev_sampled_code = cache.prev_hiddens if exists(cache) else ((None,) * 4)
+
+        # getting proposed action for the two phases
+
+        next_action_proposer_hidden = None
+
+        meta_embed = self.model_to_meta(residual_stream)
+
+        if discovery_phase:
+            mask = maybe(lens_to_mask)(episode_lens, meta_embed.shape[1])
+
+            encoded_temporal = self.bidirectional_temporal_encoder(meta_embed, mask = mask)
+
+            proposed_action_hidden, _ = self.emitter(cat((encoded_temporal, meta_embed), dim = -1))
+            to_logits = self.emitter_to_binary_logits
+
+        else: # else internal rl phase
+
+            proposed_action_hidden, next_action_proposer_hidden = self.action_proposer(meta_embed, prev_action_proposer_hidden)
+            to_logits = self.proposer_to_binary_logits
+
+        # sample from the binary mapper
+
+        binary_logits = to_logits(proposed_action_hidden)
+
+        one_hot, kl_loss = self.binary_mapper(
+            binary_logits,
+            temperature = temperature,
+            reduce_aux_kl_loss = False
+        )
+
+        # bottled action is now the one-hot sparse codes (with straight-through)
+
+        sampled_codes = one_hot
+
+        # switching unit timer
+
+        batch, seq_len, dim = sampled_codes.shape
+
+        if not exists(prev_sampled_code):
+            prev_sampled_code = torch.zeros(batch, 1, self.num_codes, device = device)
+
+        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'
+            z_prev = prev_sampled_code
+
+        switch_input = torch.cat((meta_embed, z_prev), dim=-1)
+
+        switching_unit_gru_out, next_switching_unit_gru_hidden = self.switching_unit(
+            switch_input, 
+            prev_switching_unit_gru_hidden
+        )
+
+        switch_beta = self.to_switching_unit_beta(switching_unit_gru_out).sigmoid()
+
+        # losses
+
+        switch_loss = self.zero
+
+        if discovery_phase:
+            # weight unreduced kl loss by switch gates
+
+            weighted_kl_loss = kl_loss * switch_beta
+            kl_loss = weighted_kl_loss.sum(dim = -1).mean()
+
+            # encourage less switching
+
+            switch_loss = switch_beta.mean()
+        else:
+            kl_loss = self.zero
+
+        # maybe hard switch, then use associative scan
+
+        if hard_switch:
+            hard_switch_beta = (switch_beta > 0.5).float()
+            switch_beta = straight_through(switch_beta, hard_switch_beta)
+
+        forget = 1. - switch_beta
+
+        # gated codes (or soft distribution)
+
+        gated_codes = self.switch_gating(switch_beta, sampled_codes * forget, prev = prev_switch_gated_hiddens)
+
+        next_switch_gated_codes = gated_codes[:, -1]
+
+        # decoder
+
+        decoder_out = self.decoder(gated_codes)
+
+        w1, w2 = self.to_hyper_network_weights(decoder_out)
+        hypernetwork_weight = einsum(w1, w2, '... i r, ... j r -> ... i j')
+
+        # generating the residual stream controlling signal
+
+        control_signal = einsum(residual_stream, hypernetwork_weight, '... d1, ... d1 d2 -> ... d1')
+
+        # returning
+
+        next_hiddens = (
+            next_action_proposer_hidden,
+            next_switching_unit_gru_hidden,
+            next_switch_gated_codes,
+            sampled_codes[:, -1:]
+        )
+
+        return control_signal, MetaControllerOutput(next_hiddens, binary_logits, sampled_codes, kl_loss, switch_loss)

{metacontroller_pytorch-0.0.25 → metacontroller_pytorch-0.0.27}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "metacontroller-pytorch"
-version = "0.0.25"
+version = "0.0.27"
 description = "Transformer Metacontroller"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -32,6 +32,7 @@ dependencies = [
     "memmap-replay-buffer>=0.0.23",
     "torch>=2.5",
     "torch-einops-utils>=0.0.16",
+    "vector-quantize-pytorch>=1.27.20",
     "x-evolution>=0.1.23",
     "x-mlps-pytorch",
     "x-transformers"

{metacontroller_pytorch-0.0.25 → metacontroller_pytorch-0.0.27}/tests/test_metacontroller.py

@@ -5,27 +5,30 @@ from pathlib import Path
 
 import torch
 from metacontroller.metacontroller import Transformer, MetaController
+from metacontroller.metacontroller_with_binary_mapper import MetaControllerWithBinaryMapper
 
 from einops import rearrange
 
+@param('use_binary_mapper_variant', (False, True))
 @param('action_discrete', (False, True))
 @param('switch_per_latent_dim', (False, True))
 @param('variable_length', (False, True))
 def test_metacontroller(
+    use_binary_mapper_variant,
     action_discrete,
     switch_per_latent_dim,
     variable_length
 ):
 
-    state = torch.randn(1, 1024, 384)
-    episode_lens = torch.tensor([512]) if variable_length else None
+    state = torch.randn(2, 128, 384)
+    episode_lens = torch.tensor([64, 64]) if variable_length else None
 
     if action_discrete:
-        actions = torch.randint(0, 4, (1, 1024))
+        actions = torch.randint(0, 4, (2, 128))
         action_embed_readout = dict(num_discrete = 4)
         assert_shape = (4,)
     else:
-        actions = torch.randn(1, 1024, 8)
+        actions = torch.randn(2, 128, 8)
         action_embed_readout = dict(num_continuous = 8)
         assert_shape = (8, 2)
 
@@ -44,16 +47,24 @@ def test_metacontroller(
 
     # discovery and internal rl phase with meta controller
 
-    meta_controller = MetaController(
-        dim_model = 512,
-        dim_meta_controller = 256,
-        dim_latent = 128,
-        switch_per_latent_dim = switch_per_latent_dim
-    )
+    if not use_binary_mapper_variant:
+        meta_controller = MetaController(
+            dim_model = 512,
+            dim_meta_controller = 256,
+            dim_latent = 128,
+            switch_per_latent_dim = switch_per_latent_dim
+        )
+    else:
+        meta_controller = MetaControllerWithBinaryMapper(
+            dim_model = 512,
+            dim_meta_controller = 256,
+            switch_per_code = switch_per_latent_dim,
+            dim_code_bits = 8, # 2**8 = 256 codes
+        )
 
     # discovery phase
 
-    (action_recon_loss, kl_loss, switch_loss) = model(state, actions, meta_controller = meta_controller, discovery_phase = True)
+    (action_recon_loss, kl_loss, switch_loss) = model(state, actions, meta_controller = meta_controller, discovery_phase = True, episode_lens = episode_lens)
     (action_recon_loss + kl_loss * 0.1 + switch_loss * 0.2).backward()
 
     # internal rl - done iteratively
@@ -66,7 +77,7 @@ def test_metacontroller(
 
         logits, cache = model(one_state, past_action_id, meta_controller = meta_controller, return_cache = True)
 
-        assert logits.shape == (1, 1, *assert_shape)
+        assert logits.shape == (2, 1, *assert_shape)
         past_action_id = model.action_readout.sample(logits)
 
     # evolutionary strategies over grpo
@@ -78,7 +89,8 @@ def test_metacontroller(
 
     meta_controller.save('./meta_controller.pt')
 
-    rehydrated_meta_controller = MetaController.init_and_load('./meta_controller.pt')
+    meta_controller_klass = meta_controller.__class__
+    rehydrated_meta_controller = meta_controller_klass.init_and_load('./meta_controller.pt')
 
     model.save('./trained.pt')