rxnn 0.2.56__tar.gz → 0.2.57__tar.gz

{rxnn-0.2.56 → rxnn-0.2.57}/PKG-INFO

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

{rxnn-0.2.56 → rxnn-0.2.57}/pyproject.toml

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

{rxnn-0.2.56 → rxnn-0.2.57}/src/rxnn/training/dataset.py

@@ -943,7 +943,7 @@ class MrlCurriculumDataset(Dataset):
         return self.get_tokenized_item(idx)
 
     def __len__(self) -> int:
-        return len(self.episodes)
+        return len(self.inputs if self.is_pre_tokenized else self.episodes)
 
     def get_subset(self, size: float, from_start: bool = False, **kwargs) -> "MRlCurriculumDataset":
         split_point = int(len(self.episodes) * ((1 - size) if not from_start else size))

{rxnn-0.2.56 → rxnn-0.2.57}/src/rxnn/training/mrl.py

@@ -36,6 +36,8 @@ class MrlConfig(TypedDict):
     freeze_embeddings: Optional[bool]
     embedding_lr: Optional[float]
     use_memory_warmup: Optional[bool]
+    debug_mode: Optional[bool]
+    debug_interval: Optional[int]
 
 
 class MrlStrategy(Enum):
@@ -109,7 +111,6 @@ class MRLTrainer:
             use_ddp: bool = False,
             use_amp: bool = False,
             dtype: torch.dtype = torch.float32,
-            debug_mode: bool = False,
     ):
         """
         Trainer for Memory Reinforcement Learning (MRL) algorithm for reactive models and Attention-Based Memory System.
@@ -140,7 +141,8 @@ class MRLTrainer:
         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)
-        self.debug_mode = debug_mode
+        self.debug_mode = config.get('debug_mode', False)
+        self.debug_interval = config.get('debug_interval', 10)
         # Internal update epochs config
         self.shared_update_epochs = config.get('update_epochs', 10)
         self.update_epochs = self.shared_update_epochs
@@ -606,7 +608,7 @@ class MRLTrainer:
             self.scaler.unscale_(self.optimizer)
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                            error_if_nonfinite=False)
-            if self.debug_mode:
+            if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients()
             # 4.5 Run scaled optimization step
             self.scaler.step(self.optimizer)
@@ -625,7 +627,7 @@ class MRLTrainer:
             # 4.4 Clip gradient norms
             torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                            error_if_nonfinite=False)
-            if self.debug_mode:
+            if self.debug_mode and self.epoch_step['train'] % self.debug_interval == 0:
                 self._log_gradients()
             # 4.5 Run scaled optimization step
             self.optimizer.step()

{rxnn-0.2.56 → rxnn-0.2.57}/src/rxnn/training/rl.py

@@ -33,10 +33,12 @@ class PPOConfig(TypedDict):
     use_distributed_advantage_norm: Optional[bool]
     clip_critic_values: Optional[bool]
     critic_value_clip: Optional[float]
+    debug_mode: Optional[bool]
+    debug_interval: Optional[int]
 
 
 class PPOAlgorithm(RlAlgorithm):
-    def __init__(self, config: Optional[PPOConfig] = None, debug_mode: bool = False):
+    def __init__(self, config: Optional[PPOConfig] = None):
         super(PPOAlgorithm, self).__init__()
 
         if config is None:
@@ -50,7 +52,9 @@ class PPOAlgorithm(RlAlgorithm):
         self.use_distributed_advantage_norm = config.get('use_distributed_advantage_norm', False)
         self.clip_critic_values = config.get('clip_critic_values', True)
         self.critic_value_clip = config.get('critic_value_clip', 20.0)
-        self.debug_mode = debug_mode
+        self.debug_mode = config.get('debug_mode', False)
+        self.debug_interval = config.get('debug_interval', 10)
+        self.debug_step = 0
 
     def critic_loss(self, values: torch.Tensor, ref_values: torch.Tensor) -> torch.Tensor:
         # Critic loss with clipped values
@@ -98,12 +102,16 @@ class PPOAlgorithm(RlAlgorithm):
         advantages = advantages.unsqueeze(-1)
 
         if self.debug_mode:
-            print(
-                f"Logits stats: min={new_logits.min().item():.4f}, max={new_logits.max().item():.4f}, mean={new_logits.mean().item():.4f}")
-            print(
-                f"Ratio stats: min={ratio.min().item():.4f}, max={ratio.max().item():.4f}, mean={ratio.mean().item():.4f}")
-            print(
-                f"Advantage stats: min={advantages.min().item():.4f}, max={advantages.max().item():.4f}, mean={advantages.mean().item():.4f}")
+            if self.debug_step != 0 and self.debug_step % self.debug_interval == 0:
+                self.debug_step = 0
+                print(
+                    f"Logits stats: min={new_logits.min().item():.4f}, max={new_logits.max().item():.4f}, mean={new_logits.mean().item():.4f}")
+                print(
+                    f"Ratio stats: min={ratio.min().item():.4f}, max={ratio.max().item():.4f}, mean={ratio.mean().item():.4f}")
+                print(
+                    f"Advantage stats: min={advantages.min().item():.4f}, max={advantages.max().item():.4f}, mean={advantages.mean().item():.4f}")
+            else:
+                self.debug_step += 1
 
         # c) Clipped surrogate loss
         surr1 = ratio * advantages