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

rxnn/training/models.py

@@ -4,6 +4,8 @@ from enum import Enum
 from typing import Literal, Iterator
 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):
@@ -37,6 +39,7 @@ class MLMTrainingModel(nn.Module):
         y = self.mlm_head(h)
         return y
 
+
 class JointTrainingModel(nn.Module):
     def __init__(
             self,
@@ -58,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,
@@ -153,18 +158,33 @@ 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, **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.Linear(embed_dim, 1)
+        self.value_head = nn.Sequential(
+            GatedLinearUnit(embed_dim, embed_dim, nn.SiLU()),
+            nn.LayerNorm(embed_dim),
+            nn.Linear(embed_dim, 1),
+            get_activation_layer(out_activation)
+        )
+        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)
@@ -175,4 +195,4 @@ class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipe
         else:
             x = x.mean(dim=1)
 
-        return self.value_head(x)
+        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,
         )
 
@@ -481,7 +489,7 @@ class MRLTrainer:
                 critic_state = smart_concat_critic_states(*state, max_length=self.critic_max_len,
                                                           pad_token_id=self.pad_token_id)
                 values = self.critic(critic_state['input_ids'], attention_mask=critic_state['attention_mask']).squeeze()
-                critic_loss = self.rl_algorithm.critic_loss(values, ref_values)
+                critic_loss = self.rl_algorithm.critic_loss(values, ref_values.detach())
             # 2.2 Run backpropagation with scaler
             self.critic_scaler.scale(critic_loss).backward()
             # 2.3 Unscale and clip gradients
@@ -495,7 +503,7 @@ class MRLTrainer:
             critic_state = smart_concat_critic_states(*state, max_length=self.critic_max_len,
                                                       pad_token_id=self.pad_token_id)
             values = self.critic(critic_state['input_ids'], attention_mask=critic_state['attention_mask']).squeeze()
-            critic_loss = self.rl_algorithm.critic_loss(values, ref_values)
+            critic_loss = self.rl_algorithm.critic_loss(values, ref_values.detach())
             # 2.2 Run backpropagation
             critic_loss.backward()
             # 2.3 Clip gradients
@@ -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/training/reward.py

@@ -42,6 +42,7 @@ class MrlRewardModel:
             running_mean_decay: float = 0.2,
             bleu_saved_weights: tuple = (0.5, 0.5),
             bleu_ref_weights: tuple = (0.5, 0.5),
+            tanh_reward_scale: bool = False,
             rewards_scale: float = 1.0,
     ):
         self.shared_embedding = shared_embedding.to(device)
@@ -71,6 +72,7 @@ class MrlRewardModel:
         self.running_mean_decay = running_mean_decay
         self.bleu_ref_weights = bleu_ref_weights
         self.bleu_saved_weights = bleu_saved_weights
+        self.tanh_reward_scale = tanh_reward_scale
         self.rewards_scale = rewards_scale
 
         self.prev_data_running_mean = None
@@ -175,6 +177,12 @@ class MrlRewardModel:
         self.prev_data_running_mean = (1 - self.running_mean_decay) * self._sequence_embedding(
             prev_data) + self.running_mean_decay * self.prev_data_running_mean
 
+    def _pre_scale_rewards(self, rewards: torch.Tensor) -> torch.Tensor:
+        if self.tanh_reward_scale:
+            return (rewards * 2) - 1  # Convert [0,1] to [-1,1]
+        else:
+            return rewards
+
     def __call__(
             self,
             generated: TokenizedDict,
@@ -204,5 +212,5 @@ class MrlRewardModel:
             cosine = self.negative_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
             sim_rewards = self.neg_bleu_factor * torch.tensor(bleu, device=self.device) + self.neg_cos_factor * cosine
 
-        rewards = (sim_rewards + self.len_factor * self.len_reward(generated) if self.reward_len else sim_rewards) * self.rewards_scale
+        rewards = self._pre_scale_rewards(sim_rewards + self.len_factor * self.len_reward(generated) if self.reward_len else sim_rewards) * self.rewards_scale
         return rewards.tolist()

rxnn/training/rl.py

@@ -21,8 +21,8 @@ class RlAlgorithm(ABC):
     def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         pass
 
-    def critic_loss(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
-        return self.critic_loss(rewards, values)
+    def critic_loss(self, values: torch.Tensor, ref_values: torch.Tensor) -> torch.Tensor:
+        return self.critic_loss(values, ref_values)
 
 
 class PPOConfig(TypedDict):
@@ -31,6 +31,8 @@ class PPOConfig(TypedDict):
     gae_gamma: Optional[float]
     entropy_coef: Optional[float]
     use_distributed_advantage_norm: Optional[bool]
+    clip_critic_values: Optional[bool]
+    critic_value_clip: Optional[float]
 
 
 class PPOAlgorithm(RlAlgorithm):
@@ -46,6 +48,14 @@ class PPOAlgorithm(RlAlgorithm):
         self.gae_gamma = config.get('gae_gamma', 0.99)
         self.entropy_coef = config.get('entropy_coef', 0.01)
         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', 10.0)
+
+    def critic_loss(self, values: torch.Tensor, ref_values: torch.Tensor) -> torch.Tensor:
+        # Critic loss with clipped values
+        if self.clip_critic_values:
+            values = torch.clamp(values, -self.critic_value_clip, self.critic_value_clip)
+        return self.critic_loss(values, ref_values)
 
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
                     old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
@@ -96,23 +106,24 @@ class PPOAlgorithm(RlAlgorithm):
 
         return policy_loss
 
-    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, next_value: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        T, B = rewards.shape
-        advantages = torch.zeros_like(rewards, device=values.device)
+    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor,
+                     last_value: torch.Tensor, dones: torch.Tensor):
+        trajectory_len, batch_size = rewards.shape
+        advantages = torch.zeros_like(rewards, device=rewards.device)
         last_advantage = 0
-        last_value = next_value.detach()
-
-        for t in reversed(range(T)):
-            if t == T - 1:
-                next_values = last_value
-            else:
-                next_values = values[t + 1]
-
-            # Mask next values if episode ended
-            next_values = next_values * ~dones[t]
-            delta = rewards[t] + self.gae_gamma * next_values - values[t]
-            advantages[t] = delta + self.gae_gamma * self.gae_lambda * last_advantage
+        next_value = last_value
+        next_done = torch.zeros(batch_size, device=dones.device)  # Last state is terminal
+        dones = dones.float()
+        for t in reversed(range(trajectory_len)):
+            # Check if next state is terminal
+            non_terminal = 1.0 - next_done
+
+            # Delta should not include next_value if next is terminal
+            delta = rewards[t] + self.gae_gamma * next_value * non_terminal - values[t]
+            advantages[t] = delta + self.gae_gamma * self.gae_lambda * non_terminal * last_advantage
             last_advantage = advantages[t]
+            next_value = values[t]
+            next_done = dones[t]
 
         returns = advantages + values
         return advantages, returns

rxnn/transformers/ff.py

@@ -66,6 +66,8 @@ def get_activation_layer(activation: str):
         return nn.SiLU()
     elif activation == 'sigmoid':
         return nn.Sigmoid()
+    elif activation == 'tanh':
+        return nn.Tanh()
     elif activation == 'linear':
         return LinearActivation()
     else:

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

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

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

@@ -16,16 +16,16 @@ 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=2KhNT7yx0AgUke4nmsFqzQKx_YYp78QvsLWYZjWeUgQ,6812
-rxnn/training/mrl.py,sha256=Aimiiqf_4p6dp5Ty9pY9VwetySBS_OFpCQlcVHVkO4Q,55124
-rxnn/training/reward.py,sha256=7MTVdNm5HnWmt6zFDi3TAYmnVSL_-24riOoY2F7z4x8,11290
-rxnn/training/rl.py,sha256=eL3C0yryiNBgl_xb-D-5dyYUtK4V4-K4t3a60x5ir28,5142
+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
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTgc,16247
-rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
+rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
 rxnn/transformers/layers.py,sha256=l0bXmhN7KOkCw0KTVLixWSo9Op4SesGabWJ4R4EQBMY,7988
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 rxnn/transformers/models.py,sha256=hey6tFN9gmLfWCZLjtl_9OcvIjGpWLI1IDeVnr5y8YM,10583
@@ -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.32.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.32.dist-info/METADATA,sha256=Ugq4LZZBakM-pUwpd0ZY0W_Ot2zUJHAxFtofIqAHzA8,25960
-rxnn-0.2.32.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.32.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,,