metacontroller-pytorch 0.0.15__py3-none-any.whl → 0.0.41__py3-none-any.whl

metacontroller/__init__.py

@@ -1 +1 @@
-from metacontroller.metacontroller import MetaController
+from metacontroller.metacontroller import MetaController, Transformer

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
 
@@ -18,7 +18,7 @@ from einops.layers.torch import Rearrange
 
 # external modules
 
-from x_transformers import Decoder
+from x_transformers import Encoder, Decoder
 from x_mlps_pytorch import Feedforwards
 from x_evolution import EvoStrategy
 
@@ -26,6 +26,9 @@ from discrete_continuous_embed_readout import Embed, Readout, EmbedAndReadout
 
 from assoc_scan import AssocScan
 
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, masked_mean, align_dims_left
+from torch_einops_utils.save_load import save_load
+
 # constants
 
 LinearNoBias = partial(Linear, bias = False)
@@ -55,42 +58,100 @@ def straight_through(src, tgt):
 
 MetaControllerOutput = namedtuple('MetaControllerOutput', (
     'prev_hiddens',
+    'input_residual_stream',
     'action_dist',
     'actions',
-    'kl_loss'
+    'switch_beta',
+    'kl_loss',
+    'switch_loss'
 ))
 
+def z_score(t, eps = 1e-8):
+    return (t - t.mean()) / (t.std() + eps)
+
+def policy_loss(
+    meta_controller,
+    state,
+    old_log_probs,
+    actions,
+    advantages,
+    mask,
+    episode_lens = None,
+    eps_clip = 0.2
+):
+    # get new log probs
+
+    action_dist = meta_controller.get_action_dist_for_internal_rl(state)
+    new_log_probs = meta_controller.log_prob(action_dist, actions)
+
+    # calculate ratio
+
+    ratio = (new_log_probs - old_log_probs).exp()
+
+    # align ratio and advantages
+
+    ratio, advantages = align_dims_left((ratio, advantages))
+
+    # ppo surrogate loss
+
+    surr1 = ratio * advantages
+    surr2 = ratio.clamp(1 - eps_clip, 1 + eps_clip) * advantages
+
+    losses = -torch.min(surr1, surr2)
+
+    # masking
+
+    if exists(episode_lens):
+        mask, episode_mask = align_dims_left((mask, lens_to_mask(episode_lens, losses.shape[1])))
+        mask = mask & episode_mask
+
+    return masked_mean(losses, mask)
+
+@save_load()
 class MetaController(Module):
     def __init__(
         self,
-        dim_latent,
+        dim_model,
         *,
+        dim_meta_controller = 256,
+        dim_latent = 128,
         switch_per_latent_dim = True,
         decoder_expansion_factor = 2.,
         decoder_depth = 1,
         hypernetwork_low_rank = 16,
-        assoc_scan_kwargs: dict = dict()
+        assoc_scan_kwargs: dict = dict(),
+        bidirectional_temporal_encoder_kwargs: dict = dict(
+            attn_dim_head = 32,
+            heads = 8
+        )
     ):
         super().__init__()
+        self.dim_model = dim_model
+        dim_meta = default(dim_meta_controller, dim_model)
+
+        # the linear that brings from model dimension 
+
+        self.model_to_meta = Linear(dim_model, dim_meta)
 
-        # there are two phases, the first (discovery ssl phase) uses acausal with some ssm i don't really believe in - let's just use a bidirectional GRU as placeholders
+        # there are two phases, the first (discovery ssl phase) uses acausal with some ssm i don't really believe in - let's just use bidirectional attention as placeholder
 
-        self.bidirectional_temporal_compressor = GRU(dim_latent, dim_latent, bidirectional = True) # revisit naming
+        self.bidirectional_temporal_encoder = Encoder(dim = dim_meta, depth = 1, **bidirectional_temporal_encoder_kwargs)
 
-        self.emitter = GRU(dim_latent * 2, dim_latent * 2)
-        self.emitter_to_action_mean_log_var = Readout(dim_latent * 2, num_continuous = dim_latent)
+        self.emitter = GRU(dim_meta * 2, dim_meta * 2)
+        self.emitter_to_action_mean_log_var = Readout(dim_meta * 2, num_continuous = dim_latent)
 
         # internal rl phase substitutes the acausal + emitter with a causal ssm
 
-        self.action_proposer = GRU(dim_latent, dim_latent)
-        self.action_proposer_mean_log_var = Readout(dim_latent, num_continuous = dim_latent)
+        self.action_proposer = GRU(dim_meta, dim_meta)
+        self.action_proposer_mean_log_var = Readout(dim_meta, num_continuous = dim_latent)
 
         # switching unit
 
         self.switch_per_latent_dim = switch_per_latent_dim
 
-        self.switching_unit = GRU(dim_latent, dim_latent)
-        self.to_switching_unit_beta = nn.Linear(dim_latent, dim_latent if switch_per_latent_dim else 1, bias = False)
+        self.dim_latent = dim_latent
+        self.switching_unit = GRU(dim_meta + dim_latent, dim_meta)
+        self.to_switching_unit_beta = nn.Linear(dim_meta, dim_latent if switch_per_latent_dim else 1, bias = False)
 
         self.switch_gating = AssocScan(**assoc_scan_kwargs)
 
@@ -104,16 +165,26 @@ class MetaController(Module):
             dim_in = dim_latent,
             dim = dim_decoder_hidden,
             depth = decoder_depth,
-            dim_out = 2 * hypernetwork_low_rank * dim_latent
+            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)
 
+    @property
+    def replay_buffer_field_dict(self):
+        return dict(
+            states = ('float', self.dim_model),
+            log_probs = ('float', self.dim_latent),
+            switch_betas = ('float', self.dim_latent if self.switch_per_latent_dim else 1),
+            latent_actions = ('float', self.dim_latent)
+        )
+
     def discovery_parameters(self):
         return [
-            *self.bidirectional_temporal_compressor.parameters(),
+            *self.model_to_meta.parameters(),
+            *self.bidirectional_temporal_encoder.parameters(),
             *self.emitter.parameters(),
             *self.emitter_to_action_mean_log_var.parameters(),
             *self.decoder.parameters(),
@@ -126,54 +197,99 @@ class MetaController(Module):
             *self.action_proposer_mean_log_var.parameters()
         ]
 
+    def get_action_dist_for_internal_rl(
+        self,
+        residual_stream
+    ):
+        meta_embed = self.model_to_meta(residual_stream)
+
+        proposed_action_hidden, _ = self.action_proposer(meta_embed)
+
+        return self.action_proposer_mean_log_var(proposed_action_hidden)
+
+    def log_prob(
+        self,
+        action_dist,
+        sampled_latent_action
+    ):
+        return self.action_proposer_mean_log_var.log_prob(action_dist, sampled_latent_action)
+
     def forward(
         self,
         residual_stream,
         cache: MetaControllerOutput | None = None,
         discovery_phase = False,
-        hard_switch = False,
-        temperature = 1.
+        hard_switch = None,
+        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 = cache.prev_hiddens if exists(cache) else ((None,) * 3)
+        prev_action_proposer_hidden, prev_switching_unit_gru_hidden, prev_switch_gated_hiddens, prev_sampled_latent_action = 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)
+
+        hard_switch = default(hard_switch, not discovery_phase) # think during internal RL phase, it needs to be a hard switch, then only the actions emitted during the switch is reinforced
+
         if discovery_phase:
             logger.warning('meta controller cache being passed back in for discovery phase, which does not make sense given bidirectional encoder')
 
-            temporal_compressed, _ = self.bidirectional_temporal_compressor(residual_stream)
-            temporal_compressed = reduce(temporal_compressed, '... (two d) -> ... d', 'mean', two = 2)
+            mask = maybe(lens_to_mask)(episode_lens, meta_embed.shape[1])
 
-            proposed_action_hidden, _ = self.emitter(cat((temporal_compressed, residual_stream), dim = -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
 
         else: # else internal rl phase
 
-            proposed_action_hidden, next_action_proposer_hidden = self.action_proposer(residual_stream, prev_action_proposer_hidden)
+            proposed_action_hidden, next_action_proposer_hidden = self.action_proposer(meta_embed, prev_action_proposer_hidden)
             readout = self.action_proposer_mean_log_var
 
         # sample from the gaussian as the action from the meta controller
 
         action_dist = readout(proposed_action_hidden)
 
-        sampled_action = readout.sample(action_dist, temperature = temperature)
+        sampled_latent_action = readout.sample(action_dist, temperature = temperature)
 
         # switching unit timer
 
-        batch, _, dim = sampled_action.shape
+        batch, seq_len, dim = sampled_latent_action.shape
+
+        # initialize prev sampled latent action to be zeros if not available (for first timestep and for discovery phase)
 
-        switching_unit_gru_out, next_switching_unit_gru_hidden = self.switching_unit(residual_stream, prev_switching_unit_gru_hidden)
+        if not exists(prev_sampled_latent_action):
+            prev_sampled_latent_action = torch.zeros(batch, 1, self.dim_latent, device = device)
+
+        if discovery_phase:
+            z_prev = cat((prev_sampled_latent_action, sampled_latent_action[:, :-1]), dim = 1)
+
+        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'
+            z_prev = prev_sampled_latent_action
+
+        # switch input is previous latent action and the embedding
+
+        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()
 
         # need to encourage normal distribution
 
-        kl_loss = self.zero
+        kl_loss = switch_loss = self.zero
 
         if discovery_phase:
             mean, log_var = action_dist.unbind(dim = -1)
@@ -188,6 +304,10 @@ class MetaController(Module):
             kl_loss = kl_loss * switch_beta
             kl_loss = kl_loss.sum(dim = -1).mean()
 
+            # encourage less switching
+
+            switch_loss = switch_beta.mean()
+
         # maybe hard switch, then use associative scan
 
         if hard_switch:
@@ -195,7 +315,7 @@ class MetaController(Module):
             switch_beta = straight_through(switch_beta, hard_switch_beta)
 
         forget = 1. - switch_beta
-        gated_action = self.switch_gating(switch_beta, sampled_action * forget, prev = prev_switch_gated_hiddens)
+        gated_action = self.switch_gating(switch_beta, sampled_latent_action * forget, prev = prev_switch_gated_hiddens)
 
         next_switch_gated_action = gated_action[:, -1]
 
@@ -208,27 +328,40 @@ class MetaController(Module):
 
         # generating the residual stream controlling signal
 
-        control_signal = einsum(gated_action, hypernetwork_weight, '... d1, ... d1 d2 -> ... d1')
-
-        modified_residual_stream = residual_stream + control_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_action
+            next_switch_gated_action,
+            sampled_latent_action[:, -1:]
         )
 
-        return modified_residual_stream, MetaControllerOutput(next_hiddens, action_dist, sampled_action, kl_loss)
+        # squeeze out the last dimension of switch_beta if single gate for all latent dimensions
+
+        if not self.switch_per_latent_dim:
+            switch_beta = rearrange(switch_beta, '... 1 -> ...')
+
+        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', (
+    'lower_body',
+    'meta_controller',
+    'upper_body'
+))
+
 TransformerOutput = namedtuple('TransformerOutput', (
     'residual_stream_latent',
     'prev_hiddens'
 ))
 
+@save_load()
 class Transformer(Module):
     def __init__(
         self,
@@ -243,7 +376,7 @@ class Transformer(Module):
         super().__init__()
 
         if isinstance(lower_body, dict):
-            lower_body = Decoder(dim = dim, **lower_body)
+            lower_body = Decoder(dim = dim, pre_norm_has_final_norm = False, **lower_body)
 
         if isinstance(upper_body, dict):
             upper_body = Decoder(dim = dim, **upper_body)
@@ -281,26 +414,38 @@ class Transformer(Module):
     def forward(
         self,
         state,
-        action_ids,
+        actions: Tensor | None = None,
         meta_controller: Module | None = None,
         cache: TransformerOutput | None = None,
         discovery_phase = False,
+        force_behavior_cloning = False,
         meta_controller_temperature = 1.,
         return_raw_action_dist = False,
         return_latents = False,
         return_cache = False,
+        episode_lens: Tensor | None = None
     ):
+        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)
+        if force_behavior_cloning:
+            assert not discovery_phase, 'discovery phase cannot be set to True if force behavioral cloning is set to True'
+            meta_controller = None
 
-        behavioral_cloning = not meta_controlling and not return_raw_action_dist
+        has_meta_controller = exists(meta_controller)
+
+        assert not (discovery_phase and not has_meta_controller), 'meta controller must be made available during discovery phase'
+
+        behavioral_cloning = force_behavior_cloning or (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
 
@@ -308,16 +453,31 @@ class Transformer(Module):
 
         # handle maybe behavioral cloning
 
-        if behavioral_cloning:
+        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:]
-            action_ids, target_action_ids = action_ids[:, :-1], action_ids[:, 1:]
+            actions, target_actions = actions[:, :-1], actions[:, 1:]
+
+            if exists(episode_lens):
+                episode_lens = (episode_lens - 1).clamp(min = 0)
 
         # transformer lower body
 
         with lower_transformer_context():
 
             state_embed = self.state_embed(state)
-            action_embed = self.action_embed(action_ids)
+
+            # handle no past action for first timestep
+
+            if exists(actions):
+                action_embed = self.action_embed(actions)
+            else:
+                action_embed = state_embed[:, 0:0] # empty action embed
+
+            if action_embed.shape[-2] == (state_embed.shape[-2] - 1):
+                action_embed = pad_at_dim(action_embed, (1, 0), dim = 1)
 
             embed = state_embed + action_embed
 
@@ -327,10 +487,12 @@ class Transformer(Module):
 
         with meta_controller_context():
 
-            if exists(meta_controller):
-                modified_residual_stream, next_meta_hiddens = meta_controller(residual_stream, cache = meta_hiddens, discovery_phase = discovery_phase, temperature = meta_controller_temperature)
+            if exists(meta_controller) and not behavioral_cloning:
+                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:
-                modified_residual_stream, next_meta_hiddens = residual_stream, None
+                control_signal, next_meta_hiddens = self.zero, None
+
+            modified_residual_stream = residual_stream + control_signal
 
         # modified residual stream sent back to transformer upper body
 
@@ -345,13 +507,22 @@ class Transformer(Module):
         # maybe return behavior cloning loss
 
         if behavioral_cloning:
+
+            loss_mask = maybe(lens_to_mask)(episode_lens, state.shape[1])
+
             state_dist_params = self.state_readout(attended)
-            state_clone_loss = self.state_readout.calculate_loss(state_dist_params, target_state)
+            state_clone_loss = self.state_readout.calculate_loss(state_dist_params, target_state, mask = loss_mask)
 
-            action_clone_loss = self.action_readout.calculate_loss(dist_params, target_action_ids)
+            action_clone_loss = self.action_readout.calculate_loss(dist_params, target_actions, mask = loss_mask)
 
             return state_clone_loss, action_clone_loss
 
+        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
+
         # returning
 
         return_one = not (return_latents or return_cache)
@@ -359,4 +530,4 @@ class Transformer(Module):
         if return_one:
             return dist_params
 
-        return dist_params, TransformerOutput(residual_stream, (next_lower_hiddens, next_meta_hiddens, next_upper_hiddens))
+        return dist_params, TransformerOutput(residual_stream, Hiddens(next_lower_hiddens, next_meta_hiddens, next_upper_hiddens))

metacontroller/metacontroller_with_binary_mapper.py

@@ -0,0 +1,315 @@
+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, align_dims_left
+from torch_einops_utils.save_load import save_load
+
+from vector_quantize_pytorch import BinaryMapper
+
+from metacontroller.metacontroller import MetaControllerOutput, policy_loss
+
+# 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()
+
+def log(t, eps = 1e-20):
+    return t.clamp_min(eps).log()
+
+# meta controller
+
+@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__()
+        self.dim_model = dim_model
+        assert not switch_per_code, 'switch_per_code is not supported for binary mapper'
+
+        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)
+
+    @property
+    def replay_buffer_field_dict(self):
+        return dict(
+            states = ('float', self.dim_model),
+            log_probs = ('float', self.dim_code_bits),
+            switch_betas = ('float', self.num_codes if self.switch_per_code else 1),
+            latent_actions = ('float', self.num_codes)
+        )
+
+    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 get_action_dist_for_internal_rl(
+        self,
+        residual_stream
+    ):
+        meta_embed = self.model_to_meta(residual_stream)
+
+        proposed_action_hidden, _ = self.action_proposer(meta_embed)
+
+        return self.proposer_to_binary_logits(proposed_action_hidden)
+
+    def log_prob(
+        self,
+        action_dist,
+        sampled_latent_action
+    ):
+        log_probs = stack((
+            F.logsigmoid(action_dist),
+            F.logsigmoid(-action_dist)
+        ), dim = -1)
+
+        indices = sampled_latent_action.argmax(dim = -1)
+        codes = self.binary_mapper.codes[indices].long()
+
+        codes = rearrange(codes, '... -> ... 1')
+        action_log_probs = log_probs.gather(-1, codes)
+        action_log_probs = rearrange(action_log_probs, '... 1 -> ...')
+
+        return action_log_probs
+
+    def forward(
+        self,
+        residual_stream,
+        cache: MetaControllerOutput | None = None,
+        discovery_phase = False,
+        hard_switch = None,
+        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)
+
+        hard_switch = default(hard_switch, not discovery_phase) # think during internal RL phase, it needs to be a hard switch, then only the actions emitted during the switch is reinforced
+
+        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
+
+            kl_loss, switch_beta = align_dims_left((kl_loss, switch_beta))
+
+            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:]
+        )
+
+        # squeeze out the last dimension of switch_beta if single gate for all codes
+
+        if not self.switch_per_code:
+            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/transformer_with_resnet.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import torch
+from torch import nn, Tensor
+from torch.nn import Module, ModuleList
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+from metacontroller.metacontroller import Transformer
+
+from torch_einops_utils import pack_with_inverse
+
+# resnet components
+
+def exists(v):
+    return v is not None
+
+class BasicBlock(Module):
+    expansion = 1
+
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        stride = 1,
+        downsample: Module | None = None
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2d(dim, dim_out, 3, stride = stride, padding = 1, bias = False)
+        self.bn1 = nn.BatchNorm2d(dim_out)
+        self.relu = nn.ReLU(inplace = True)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, padding = 1, bias = False)
+        self.bn2 = nn.BatchNorm2d(dim_out)
+        self.downsample = downsample
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if exists(self.downsample):
+            identity = self.downsample(x)
+
+        out += identity
+        return self.relu(out)
+
+class Bottleneck(Module):
+    expansion = 4
+
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        stride = 1,
+        downsample: Module | None = None
+    ):
+        super().__init__()
+        width = dim_out # simple resnet shortcut
+        self.conv1 = nn.Conv2d(dim, width, 1, bias = False)
+        self.bn1 = nn.BatchNorm2d(width)
+        self.conv2 = nn.Conv2d(width, width, 3, stride = stride, padding = 1, bias = False)
+        self.bn2 = nn.BatchNorm2d(width)
+        self.conv3 = nn.Conv2d(width, dim_out * self.expansion, 1, bias = False)
+        self.bn3 = nn.BatchNorm2d(dim_out * self.expansion)
+        self.relu = nn.ReLU(inplace = True)
+        self.downsample = downsample
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if exists(self.downsample):
+            identity = self.downsample(x)
+
+        out += identity
+        return self.relu(out)
+
+class ResNet(Module):
+    def __init__(
+        self,
+        block: type[BasicBlock | Bottleneck],
+        layers: list[int],
+        num_classes = 1000,
+        channels = 3
+    ):
+        super().__init__()
+        self.inplanes = 64
+
+        self.conv1 = nn.Conv2d(channels, 64, kernel_size = 7, stride = 2, padding = 3, bias = False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace = True)
+        self.maxpool = nn.MaxPool2d(kernel_size = 3, stride = 2, padding = 1)
+
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride = 2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride = 2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride = 2)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.flatten = Rearrange('b c 1 1 -> b c')
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(
+        self,
+        block: type[BasicBlock | Bottleneck],
+        planes: int,
+        blocks: int,
+        stride: int = 1
+    ) -> nn.Sequential:
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride = stride, bias = False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+# resnet factory
+
+def resnet18(num_classes: any = 1000):
+    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes)
+
+def resnet34(num_classes: any = 1000):
+    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes)
+
+def resnet50(num_classes: any = 1000):
+    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes)
+
+# transformer with resnet
+
+class TransformerWithResnet(Transformer):
+    def __init__(
+        self,
+        *,
+        resnet_type = 'resnet18',
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        resnet_klass = resnet18
+        if resnet_type == 'resnet34':
+            resnet_klass = resnet34
+        elif resnet_type == 'resnet50':
+            resnet_klass = resnet50
+
+        self.resnet_dim = kwargs['state_embed_readout']['num_continuous']
+        self.visual_encoder = resnet_klass(num_classes = self.resnet_dim)
+
+    def visual_encode(self, x: Tensor) -> Tensor:
+        if x.shape[-1] == 3:
+            x = rearrange(x, '... h w c -> ... c h w')
+
+        x, inverse = pack_with_inverse(x, '* c h w')
+
+        h = self.visual_encoder(x)
+
+        return inverse(h, '* d')

{metacontroller_pytorch-0.0.15.dist-info → metacontroller_pytorch-0.0.41.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.15
+Version: 0.0.41
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -39,7 +39,10 @@ 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.25
+Requires-Dist: torch-einops-utils>=0.0.19
 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
@@ -54,13 +57,110 @@ Description-Content-Type: text/markdown
 
 Implementation of the MetaController proposed in [Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning](https://arxiv.org/abs/2512.20605)
 
+## Install
+
+```shell
+$ pip install metacontroller-pytorch
+```
+
+## Appreciation
+
+- [Pranoy](https://github.com/pranoyr) for submitting a pull request for fixing the previous latent action not being included in the inputs to the switching unit
+
+- [Diego Calanzone](https://github.com/ddidacus) for proposing testing on BabyAI gridworld task, and submitting the [pull request](https://github.com/lucidrains/metacontroller/pull/3) for behavior cloning and discovery phase training for it!
+
+## Usage
+
+```python
+import torch
+from metacontroller import Transformer, MetaController
+
+# 1. initialize model
+
+model = Transformer(
+    dim = 512,
+    action_embed_readout = dict(num_discrete = 4),
+    state_embed_readout = dict(num_continuous = 384),
+    lower_body = dict(depth = 2),
+    upper_body = dict(depth = 2)
+)
+
+state = torch.randn(2, 128, 384)
+actions = torch.randint(0, 4, (2, 128))
+
+# 2. behavioral cloning (BC)
+
+state_loss, action_loss = model(state, actions)
+(state_loss + action_loss).backward()
+
+# 3. discovery phase
+
+meta_controller = MetaController(
+    dim_model = 512,
+    dim_meta_controller = 256,
+    dim_latent = 128
+)
+
+action_recon_loss, kl_loss, switch_loss = model(
+    state,
+    actions,
+    meta_controller = meta_controller,
+    discovery_phase = True
+)
+
+(action_recon_loss + kl_loss + switch_loss).backward()
+
+# 4. internal rl phase (GRPO)
+
+# ... collect trajectories ...
+
+logits, cache = model(
+    one_state,
+    past_action_id,
+    meta_controller = meta_controller,
+    return_cache = True
+)
+
+meta_output = cache.prev_hiddens.meta_controller
+old_log_probs = meta_controller.log_prob(meta_output.action_dist, meta_output.actions)
+
+# ... calculate advantages ...
+
+loss = meta_controller.policy_loss(
+    group_states,
+    group_old_log_probs,
+    group_latent_actions,
+    group_advantages,
+    group_switch_betas
+)
+
+loss.backward()
+```
+
+Or using [evolutionary strategies](https://arxiv.org/abs/2511.16652) for the last portion
+
+```python
+# 5. evolve (ES over GRPO)
+
+model.meta_controller = meta_controller
+
+def environment_callable(model):
+    # return a fitness score
+    return 1.0
+
+model.evolve(
+    num_generations = 10,
+    environment = environment_callable
+)
+```
+
 ## Citations
 
 ```bibtex
 @misc{kobayashi2025emergenttemporalabstractionsautoregressive,
     title   = {Emergent temporal abstractions in autoregressive models enable hierarchical reinforcement learning}, 
     author  = {Seijin Kobayashi and Yanick Schimpf and Maximilian Schlegel and Angelika Steger and Maciej Wolczyk and Johannes von Oswald and Nino Scherrer and Kaitlin Maile and Guillaume Lajoie and Blake A. Richards and Rif A. Saurous and James Manyika and Blaise Agüera y Arcas and Alexander Meulemans and João Sacramento},
-    year={2025},
+    year    = {2025},
     eprint  = {2512.20605},
     archivePrefix = {arXiv},
     primaryClass = {cs.LG},
@@ -78,3 +178,29 @@ Implementation of the MetaController proposed in [Emergent temporal abstractions
     url     = {https://api.semanticscholar.org/CorpusID:279464702}
 }
 ```
+
+```bibtex
+@misc{hwang2025dynamicchunkingendtoendhierarchical,
+    title   = {Dynamic Chunking for End-to-End Hierarchical Sequence Modeling},
+    author  = {Sukjun Hwang and Brandon Wang and Albert Gu},
+    year    = {2025},
+    eprint  = {2507.07955},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2507.07955},
+}
+```
+
+```bibtex
+@misc{fleuret2025freetransformer,
+    title     = {The Free Transformer}, 
+    author    = {François Fleuret},
+    year      = {2025},
+    eprint    = {2510.17558},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url       = {https://arxiv.org/abs/2510.17558}, 
+}
+```
+
+*Life can only be understood backwards; but it must be lived forwards* - Søren Kierkegaard

metacontroller_pytorch-0.0.41.dist-info/RECORD

@@ -0,0 +1,8 @@
+metacontroller/__init__.py,sha256=iSKbCDp3UrWhZg7SIJFYNjdVQU56u-vqZarE6qCSX74,70
+metacontroller/metacontroller.py,sha256=bhgCqqM-dfysGrMtZYe2w87lRVkf8fETjxUCdjrnI8Q,17386
+metacontroller/metacontroller_with_binary_mapper.py,sha256=Ce5-O95_pLuWNA3aZTlKrTGbc5cemb61tBtJBdSiLx4,9843
+metacontroller/transformer_with_resnet.py,sha256=R49ycusbq3kEX97WHZ41WY2ONc2mYPOuRUCmaFcBOEo,5546
+metacontroller_pytorch-0.0.41.dist-info/METADATA,sha256=IvP-wC73xCnT8X1aul1IfcaC4fUwRq9Y2UB1h0JG5TI,6822
+metacontroller_pytorch-0.0.41.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+metacontroller_pytorch-0.0.41.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+metacontroller_pytorch-0.0.41.dist-info/RECORD,,

metacontroller_pytorch-0.0.15.dist-info/RECORD

@@ -1,6 +0,0 @@
-metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
-metacontroller/metacontroller.py,sha256=ug3xeMTZKApTF8oOPx9hWypeDjRflf1IJp8RiysXgTo,11618
-metacontroller_pytorch-0.0.15.dist-info/METADATA,sha256=9d39BpcuVeOVVSD66lCVHCK1GjrkeKzRtxKOPOc-7xQ,3736
-metacontroller_pytorch-0.0.15.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-metacontroller_pytorch-0.0.15.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-metacontroller_pytorch-0.0.15.dist-info/RECORD,,