metacontroller-pytorch 0.0.29__tar.gz → 0.0.31__tar.gz

{metacontroller_pytorch-0.0.29 → metacontroller_pytorch-0.0.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.29
+Version: 0.0.31
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller

{metacontroller_pytorch-0.0.29 → metacontroller_pytorch-0.0.31}/metacontroller/metacontroller.py

@@ -58,8 +58,10 @@ def straight_through(src, tgt):
 
 MetaControllerOutput = namedtuple('MetaControllerOutput', (
     'prev_hiddens',
+    'input_residual_stream',
     'action_dist',
     'actions',
+    'switch_beta',
     'kl_loss',
     'switch_loss'
 ))
@@ -150,7 +152,7 @@ class MetaController(Module):
         residual_stream,
         cache: MetaControllerOutput | None = None,
         discovery_phase = False,
-        hard_switch = False,
+        hard_switch = None,
         temperature = 1.,
         episode_lens: Tensor | None = None
     ):
@@ -166,6 +168,8 @@ class MetaController(Module):
 
         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')
 
@@ -268,7 +272,7 @@ class MetaController(Module):
             sampled_latent_action[:, -1:]
         )
 
-        return control_signal, MetaControllerOutput(next_hiddens, action_dist, sampled_latent_action, kl_loss, switch_loss)
+        return control_signal, MetaControllerOutput(next_hiddens, residual_stream, action_dist, sampled_latent_action, switch_beta, kl_loss, switch_loss)
 
 # main transformer, which is subsumed into the environment after behavioral cloning
 

{metacontroller_pytorch-0.0.29 → metacontroller_pytorch-0.0.31}/metacontroller/metacontroller_with_binary_mapper.py

@@ -28,6 +28,8 @@ from torch_einops_utils.save_load import save_load
 
 from vector_quantize_pytorch import BinaryMapper
 
+from metacontroller.metacontroller import MetaControllerOutput
+
 # constants
 
 LinearNoBias = partial(Linear, bias = False)
@@ -50,14 +52,6 @@ def straight_through(src, tgt):
 
 # meta controller
 
-MetaControllerOutput = namedtuple('MetaControllerOutput', (
-    'prev_hiddens',
-    'action_dist',
-    'codes',
-    'kl_loss',
-    'switch_loss'
-))
-
 @save_load()
 class MetaControllerWithBinaryMapper(Module):
     def __init__(
@@ -148,7 +142,7 @@ class MetaControllerWithBinaryMapper(Module):
         residual_stream,
         cache: MetaControllerOutput | None = None,
         discovery_phase = False,
-        hard_switch = False,
+        hard_switch = None,
         temperature = 1.,
         episode_lens: Tensor | None = None
     ):
@@ -164,6 +158,8 @@ class MetaControllerWithBinaryMapper(Module):
 
         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])
 
@@ -265,4 +261,4 @@ class MetaControllerWithBinaryMapper(Module):
             sampled_codes[:, -1:]
         )
 
-        return control_signal, MetaControllerOutput(next_hiddens, binary_logits, sampled_codes, kl_loss, switch_loss)
+        return control_signal, MetaControllerOutput(next_hiddens, residual_stream, binary_logits, sampled_codes, switch_beta, kl_loss, switch_loss)

{metacontroller_pytorch-0.0.29 → metacontroller_pytorch-0.0.31}/pyproject.toml

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

{metacontroller_pytorch-0.0.29 → metacontroller_pytorch-0.0.31}/tests/test_metacontroller.py

@@ -59,7 +59,7 @@ def test_metacontroller(
             dim_model = 512,
             dim_meta_controller = 256,
             switch_per_code = switch_per_latent_dim,
-            dim_code_bits = 8, # 2**8 = 256 codes
+            dim_code_bits = 8, # 2 ** 8 = 256 codes
         )
 
     # discovery phase