rxnn 0.2.47__tar.gz → 0.2.48__tar.gz

{rxnn-0.2.47 → rxnn-0.2.48}/PKG-INFO

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

{rxnn-0.2.47 → rxnn-0.2.48}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.47"
+version = "0.2.48"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.47 → rxnn-0.2.48}/src/rxnn/memory/norm.py

@@ -163,7 +163,7 @@ def init_memory_norm(
     init_scale: float = 1.0,
     per_dim_scale: bool = False,
 ) -> nn.Module:
-    assert norm_type in ['layer', 'rms', 'adaptive', 'positional']
+    assert norm_type in ['layer', 'rms', 'adaptive', 'positional', 'classic-rms']
     if norm_type == 'layer':
         return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)
     elif norm_type == 'rms':
@@ -172,4 +172,6 @@ def init_memory_norm(
         return AdaptiveRMSMemoryNorm(dim, use_gate, decay, init_scale, init_gate)
     elif norm_type == 'positional':
         return AdaptivePositionalMemoryNorm(num_slots, dim, decay, use_scale, use_gate, init_gate, per_dim_scale)
+    elif norm_type == 'classic-rms':
+        return nn.RMSNorm(dim)
     return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)

{rxnn-0.2.47 → rxnn-0.2.48}/src/rxnn/training/mrl.py

@@ -35,6 +35,7 @@ class MrlConfig(TypedDict):
     moe_aux_loss_scale: Optional[float]
     freeze_embeddings: Optional[bool]
     embedding_lr: Optional[float]
+    use_memory_warmup: Optional[bool]
 
 
 class MrlStrategy(Enum):
@@ -136,6 +137,7 @@ class MRLTrainer:
         self.moe_aux_loss_scale = config.get('moe_aux_loss_scale', 0.01)
         self.shared_freeze_embeddings = config.get('freeze_embeddings', False)
         self.freeze_embeddings = self.shared_freeze_embeddings
+        self.use_memory_warmup = config.get('use_memory_warmup', False)
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -381,6 +383,11 @@ class MRLTrainer:
             self.writer.add_scalar(f'Collect/episode reward (steps: {self.curriculum_steps})', avg_reward,
                                    self.stage_step['collect'])
 
+    def memory_warmup(self, query: TokenizedDict, answer: TokenizedDict):
+        if self.use_memory_warmup:
+            with torch.no_grad():
+                self.encode_and_update_stm(query, answer)
+
     def collect_trajectories(self, dataloader: DataLoader, epoch: int, batch_size: int) -> list[MrlTrajectoryEpisode]:
         """Collect trajectories for PPO for current curriculum step."""
         # 1. Init trajectories list
@@ -402,8 +409,13 @@ class MRLTrainer:
                     first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
                     interactions = interactions[:self.curriculum_steps]
                     interactions_len = len(interactions)
+
+                    first_interaction = self._move_multiple_batches(first_query, first_answer)
+
+                    if reset_done:
+                        self.memory_warmup(*first_interaction)
                     # 6. Encode and update STM with data to save from first interaction
-                    self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
+                    self.encode_and_update_stm(*first_interaction)
 
                     # 7. Save first interaction as data to save (for trajectory state)
                     query, answer = first_query, first_answer
@@ -649,6 +661,9 @@ class MRLTrainer:
 
                 self.actor.clone_reset_memory()
 
+                if should_reset_stm and step_idx == 0:
+                    self.memory_warmup(query, answer)
+
                 # 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)
@@ -798,13 +813,16 @@ class MRLTrainer:
                 if batch['query']['input_ids'].size(0) == batch_size:
                     self._increment_steps('eval')
                     # 3. Reset STM with random resets ratio and reward model running mean
-                    self.reset_stm()
+                    reset_stm = self.reset_stm()
                     self.reward.reset_running_mean()
 
                     # 4. Get batches for first queries, answers and all follow-up interactions
                     first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
                     # 5. Encode and update STM with initial interactions (batch)
-                    self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
+                    first_interaction = self._move_multiple_batches(first_query, first_answer)
+                    if reset_stm:
+                        self.memory_warmup(*first_interaction)
+                    self.encode_and_update_stm(*first_interaction)
 
                     # 6. Save follow-up interactions len and first query and answer as previous one for iteration
                     interactions_len = len(interactions)