rxnn 0.2.33__py3-none-any.whl → 0.2.34__py3-none-any.whl

rxnn/training/models.py

@@ -6,6 +6,7 @@ from huggingface_hub import PyTorchModelHubMixin
 from ..transformers.models import ReactiveTransformerEncoder, ReactiveTransformerDecoder
 from ..transformers.ff import GatedLinearUnit, get_activation_layer
 
+
 class MLMHead(nn.Module, PyTorchModelHubMixin, license="apache-2.0"):
     def __init__(self, embed_dim: int, vocab_size: int, *args, **kwargs):
         super(MLMHead, self).__init__(*args, **kwargs)
@@ -38,6 +39,7 @@ class MLMTrainingModel(nn.Module):
         y = self.mlm_head(h)
         return y
 
+
 class JointTrainingModel(nn.Module):
     def __init__(
             self,
@@ -59,10 +61,12 @@ class JointTrainingModel(nn.Module):
         y_d = self.decoder(x_d, attention_mask=attention_mask)
         return y_e, y_d
 
+
 class MrlActorAction(Enum):
     DECODE = 1
     UPDATE = 2
 
+
 class MrlActorModel(nn.Module):
     def __init__(
             self,
@@ -154,15 +158,18 @@ class MrlActorModel(nn.Module):
             list(self.memory_attention.parameters())
         ))
 
-    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None, action: MrlActorAction = MrlActorAction.DECODE) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None,
+                action: MrlActorAction = MrlActorAction.DECODE) -> torch.Tensor:
         if action == MrlActorAction.DECODE:
             return self.decoder(x, attention_mask=attention_mask)
         else:
             _, ed = self.encoder(x, attention_mask=attention_mask)
             return self.memory_attention(ed, attention_mask=attention_mask)
 
+
 class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipeline_tag="text-classification"):
-    def __init__(self, encoder: nn.Module, embed_dim: int, out_activation: Literal['sigmoid', 'tanh', 'linear'] = 'sigmoid', output_scale: float = 1.0,  **kwargs):
+    def __init__(self, encoder: nn.Module, embed_dim: int,
+                 out_activation: Literal['sigmoid', 'tanh', 'linear'] = 'sigmoid', output_scale: float = 1.0, **kwargs):
         super(MrlCriticModel, self).__init__(**kwargs)
         self.encoder = encoder
         self.value_head = nn.Sequential(
@@ -173,6 +180,12 @@ class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipe
         )
         self.output_scale = output_scale
 
+    def head_parameters(self) -> Iterator[nn.Parameter]:
+        return self.value_head.parameters()
+
+    def encoder_parameters(self) -> Iterator[nn.Parameter]:
+        return self.encoder.parameters()
+
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         x, _ = self.encoder(x, attention_mask=attention_mask)
 
@@ -183,4 +196,3 @@ class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipe
             x = x.mean(dim=1)
 
         return self.value_head(x) * self.output_scale
-

rxnn/training/mrl.py

@@ -15,11 +15,13 @@ from .reward import MrlRewardMode, MrlRewardModel
 from .models import MrlActorAction, MrlActorModel, MrlCriticModel
 from .ddp import get_os_ddp_config, distributed_mean
 
+
 class MrlConfig(TypedDict):
     lr: float
     separate_memory_lr: Optional[bool]
     memory_lr: Optional[float]
     critic_lr: float
+    critic_encoder_lr: float
     max_seq_len: int
     critic_max_len: int
     weight_decay: float
@@ -58,6 +60,7 @@ class CurriculumConfig(TypedDict):
     lr: Optional[float]
     memory_lr: Optional[float]
     critic_lr: Optional[float]
+    critic_encoder_lr: Optional[float]
     weight_decay: Optional[float]
     critic_weight_decay: Optional[float]
     update_epochs: Optional[int]
@@ -158,6 +161,7 @@ class MRLTrainer:
                 'critic_lr': config.get('critic_lr', 1e-4),
                 'weight_decay': config.get('weight_decay', 0.01),
                 'critic_weight_decay': config.get('critic_weight_decay', 0.01),
+                'critic_encoder_lr': config.get('critic_encoder_lr', config.get('critic_lr', 1e-4)),
             }
         else:
             self.base_optim_config = {
@@ -165,6 +169,7 @@ class MRLTrainer:
                 'critic_lr': config.get('critic_lr', 1e-4),
                 'weight_decay': config.get('weight_decay', 0.01),
                 'critic_weight_decay': config.get('critic_weight_decay', 0.01),
+                'critic_encoder_lr': config.get('critic_encoder_lr', config.get('critic_lr', 1e-4)),
             }
 
         self.optim_config = self.base_optim_config
@@ -202,6 +207,7 @@ class MRLTrainer:
             critic_lr: float,
             weight_decay: float,
             critic_weight_decay: float,
+            critic_encoder_lr: float,
             memory_lr: Optional[float] = None,
     ) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
         if memory_lr is not None:
@@ -219,8 +225,10 @@ class MRLTrainer:
             )
 
         critic_optimizer = torch.optim.AdamW(
-            self.critic.parameters(),
-            lr=critic_lr,
+            [
+                {'params': self.critic.head_parameters(), 'lr': critic_lr},
+                {'params': self.critic.encoder_parameters(), 'lr': critic_encoder_lr},
+            ],
             weight_decay=critic_weight_decay,
         )
 
@@ -633,7 +641,8 @@ class MRLTrainer:
                 for i, t in enumerate(episode['steps'])
             ]
             values = torch.stack([
-                self._critic_values_with_memory(r, *self._move_multiple_batches(*t['state'])) for t, r in flat_trajectories
+                self._critic_values_with_memory(r, *self._move_multiple_batches(*t['state'])) for t, r in
+                flat_trajectories
             ]).to(self.device)
             rewards = torch.stack([torch.tensor(t['reward']) for t, _ in flat_trajectories]).to(self.device)
             dones = torch.stack([torch.tensor(t['done']) for t, _ in flat_trajectories]).to(self.device)
@@ -646,7 +655,8 @@ class MRLTrainer:
             dones = torch.stack([torch.tensor(t['done']) for t in flat_trajectories]).to(self.device)
         return values, rewards, dones
 
-    def _critic_values_with_memory(self, reset_stm: bool, *moved_state: tuple[TokenizedDict, TokenizedDict, TokenizedDict]) -> torch.Tensor:
+    def _critic_values_with_memory(self, reset_stm: bool,
+                                   *moved_state: tuple[TokenizedDict, TokenizedDict, TokenizedDict]) -> torch.Tensor:
         # 1. Calculate critic values in memory aware version - reset/update STM before calculating values
         with torch.no_grad():
             # 2. Reset STM if it was reset in trajectory collection
@@ -933,6 +943,7 @@ class MRLTrainer:
                     'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
                     'critic_weight_decay': config.get('critic_weight_decay',
                                                       self.base_optim_config['critic_weight_decay']),
+                    'critic_encoder_lr': config.get('critic_encoder_lr', self.base_optim_config['critic_encoder_lr']),
                     'memory_lr': config.get('memory_lr', self.base_optim_config['memory_lr']),
                 }
             else:
@@ -942,6 +953,7 @@ class MRLTrainer:
                     'weight_decay': config.get('weight_decay', self.base_optim_config['weight_decay']),
                     'critic_weight_decay': config.get('critic_weight_decay',
                                                       self.base_optim_config['critic_weight_decay']),
+                    'critic_encoder_lr': config.get('critic_encoder_lr', self.base_optim_config['critic_encoder_lr']),
                 }
             self.optimizer, self.critic_optimizer = self._init_optimizers(**self.optim_config)
         elif self.optim_config != self.base_optim_config:

{rxnn-0.2.33.dist-info → rxnn-0.2.34.dist-info}/METADATA

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

{rxnn-0.2.33.dist-info → rxnn-0.2.34.dist-info}/RECORD

@@ -16,8 +16,8 @@ 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=8FV5eZx1HxtqRSgikwfKoB_bNhPuMYyNi0uSXB65-M4,7223
-rxnn/training/mrl.py,sha256=1pYzjXI17FDZGPTVpmbaBvMYpB-a6SLv-84RHXA4JEA,55142
+rxnn/training/models.py,sha256=y-9XHedSheyK1AmLBp3ayulnUvAmDuJ3t0qVg8wHBRg,7463
+rxnn/training/mrl.py,sha256=fIrg1Er0aAK4TnyDRmJC1m7az9wdkhikxv0CBCrGT-c,55868
 rxnn/training/reward.py,sha256=B7nerPk9eNAv2i7umtNF88tVQVwijNNrchIrEITGHKk,11623
 rxnn/training/rl.py,sha256=ckx1nlzIGZBabzwZNRj4isvHqRZwg0y0jGOT-SN6KZc,5841
 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.33.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.33.dist-info/METADATA,sha256=im17irb58IYMXOzMXE6QaSPF31Akx0iYS4ay-aRqA9Q,25960
-rxnn-0.2.33.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.33.dist-info/RECORD,,
+rxnn-0.2.34.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.34.dist-info/METADATA,sha256=Q7LqPr7KHFhMPL6UrbqG1SmtJbM2Ho-Yuxp_7LyCtYw,25960
+rxnn-0.2.34.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.34.dist-info/RECORD,,