rxnn 0.2.43__py3-none-any.whl → 0.2.44__py3-none-any.whl

rxnn/memory/stm.py

@@ -62,6 +62,9 @@ class ShortTermMemory(nn.Module):
     def reset(self, init_type: str = None):
         self.memory = self._init_tensor(init_type).to(self.memory.device)
 
+    def clone_detach_reset(self):
+        self.memory = self.memory.detach().clone()
+
     def resize(self, new_stm_size: int, init_type: str = None):
         self.stm_size = new_stm_size
         delattr(self, 'memory')

rxnn/rxt/models.py

@@ -301,6 +301,9 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
     def reset_memory(self, init_type: str = None):
         self.model.stm.reset(init_type)
 
+    def clone_reset_memory(self):
+        self.model.stm.clone_detach_reset()
+
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)
 

rxnn/training/models.py

@@ -146,6 +146,9 @@ class MrlActorModel(nn.Module):
     def reset_memory(self):
         self.memory_attention.reset_memory()
 
+    def clone_reset_memory(self):
+        self.memory_attention.clone_reset_memory()
+
     def memory_parameters(self) -> list[nn.Parameter]:
         return list(set(
             self.encoder.memory_parameters() +

rxnn/training/mrl.py

@@ -646,6 +646,8 @@ class MRLTrainer:
                 step_critic_values = episode_critic_values[step_idx]
                 step_advantages = episode_advantages[step_idx]
 
+                self.actor.clone_reset_memory()
+
                 # 7. In memory aware critic version, encode and update STM before critic update, to include its gradients in critic loss too
                 if self.memory_aware_critic:
                     self.encode_and_update_stm(query, answer)

{rxnn-0.2.43.dist-info → rxnn-0.2.44.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.43
+Version: 0.2.44
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.43.dist-info → rxnn-0.2.44.dist-info}/RECORD

@@ -7,17 +7,17 @@ rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/attention.py,sha256=POszZeW0QBKOh4VTDVekmZGKKwUr1Zj0FOAilTv8Vyg,2411
 rxnn/memory/norm.py,sha256=E98jOQEuIOFFhlkvS8s4fFN-D4tLO6vaOqnObv1oVmA,6592
-rxnn/memory/stm.py,sha256=SSfc-RL9FE-RLkmOEkLB-9Rb00ZXbMLbsAEPdpIW89o,3851
+rxnn/memory/stm.py,sha256=jv57gsH9XW19sLbxpRDqsp1yfsii_4Ef4Ncr_ztk-i4,3937
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=lRn7NRIAAeCxr8hoIXanhaD-cGwVwA23hBdIQpBK6kc,14484
+rxnn/rxt/models.py,sha256=jh7TNLu_7CL0PH_T99rMZHcLezFPiZi-xnPazNyn_dU,14563
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
 rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
 rxnn/training/callbacks.py,sha256=p72lbzFAmFjpcUvyy4aUB3qd53I8C6Sk5w9nQvsKgTk,35852
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
 rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
-rxnn/training/models.py,sha256=7IcUsqMdQ3NldXAxHf9mlZw0rmCfSFlRK51nGe4qLAg,8996
-rxnn/training/mrl.py,sha256=4_FYwfI71adrbmDLq9TaBTgpiU8lLAjuR5QYMd57WN4,59423
+rxnn/training/models.py,sha256=tqABOt_xEcWbZNEW2I2Jt-3eyaGICK011zILwuTk6Zc,9082
+rxnn/training/mrl.py,sha256=BvrwqrIIyg_EmUA5p7c6UBcfFQ0ePIcl-EHEFQqyl2E,59472
 rxnn/training/reward.py,sha256=B7nerPk9eNAv2i7umtNF88tVQVwijNNrchIrEITGHKk,11623
 rxnn/training/rl.py,sha256=q4NzIZAmXRHVToT13IHrPTtEikWQUvT0NO0IjApjAO8,6171
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -33,7 +33,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.43.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.43.dist-info/METADATA,sha256=5k76J0DHE9dpJAteFEPaJO8TYp-D8DrQpD26UeUerJw,25960
-rxnn-0.2.43.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.43.dist-info/RECORD,,
+rxnn-0.2.44.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.44.dist-info/METADATA,sha256=tW2Ve4whRK2LfCxix10dTLS5Dl_0C6KhcK8FsoKq-x0,25960
+rxnn-0.2.44.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.44.dist-info/RECORD,,