metacontroller-pytorch 0.0.35__py3-none-any.whl → 0.0.37__py3-none-any.whl

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', (
@@ -406,6 +408,7 @@ class Transformer(Module):
         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,
@@ -414,17 +417,25 @@ 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)
+        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
+
+        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 = 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
 
@@ -432,7 +443,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:]
@@ -465,7 +477,7 @@ class Transformer(Module):
 
         with meta_controller_context():
 
-            if exists(meta_controller):
+            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:
                 control_signal, next_meta_hiddens = self.zero, None
@@ -495,7 +507,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/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.dist-info → metacontroller_pytorch-0.0.37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.35
+Version: 0.0.37
 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.37.dist-info/RECORD

@@ -0,0 +1,8 @@
+metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
+metacontroller/metacontroller.py,sha256=3cVg4TkfD8bkuED0mcGcfAEjJujcJ9tf_qMB8ict12c,17017
+metacontroller/metacontroller_with_binary_mapper.py,sha256=odZs49ZWY7_FfEweYkD0moX7Vn0jGd91FjFTxzjLyr8,9480
+metacontroller/transformer_with_resnet.py,sha256=R49ycusbq3kEX97WHZ41WY2ONc2mYPOuRUCmaFcBOEo,5546
+metacontroller_pytorch-0.0.37.dist-info/METADATA,sha256=4mkDBWI-ma5TR38PpAIxEKj6VlVlQOYvJQAGVswQ3IQ,4747
+metacontroller_pytorch-0.0.37.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+metacontroller_pytorch-0.0.37.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+metacontroller_pytorch-0.0.37.dist-info/RECORD,,

metacontroller_pytorch-0.0.35.dist-info/RECORD

@@ -1,8 +0,0 @@
-metacontroller/__init__.py,sha256=lj7IOGpN--qMxJWbB-4SGqoPXG7Hd4mgtToTRSyTZ58,57
-metacontroller/metacontroller.py,sha256=sj-cHpYm9NHZEBKbLQaf4MtZCv2lcBI2cAyj5Y9bAgc,16410
-metacontroller/metacontroller_with_binary_mapper.py,sha256=9mMKMp3zVQzjbJvoC1dBRibarHHgjnOf1tRyeY1VvAM,9423
-metacontroller/transformer_with_resnet.py,sha256=R49ycusbq3kEX97WHZ41WY2ONc2mYPOuRUCmaFcBOEo,5546
-metacontroller_pytorch-0.0.35.dist-info/METADATA,sha256=eV2Y0yW-iY2_I0gPyCA8OqChqVWFwh3GkJZFzQcZ2a0,4747
-metacontroller_pytorch-0.0.35.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-metacontroller_pytorch-0.0.35.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-metacontroller_pytorch-0.0.35.dist-info/RECORD,,