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

rxnn/training/models.py

@@ -4,6 +4,7 @@ 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):
@@ -161,10 +162,16 @@ class MrlActorModel(nn.Module):
             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 forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         x, _ = self.encoder(x, attention_mask=attention_mask)
@@ -175,4 +182,5 @@ 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

@@ -481,7 +481,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 +495,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

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.33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.32
+Version: 0.2.33
 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.33.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=8FV5eZx1HxtqRSgikwfKoB_bNhPuMYyNi0uSXB65-M4,7223
+rxnn/training/mrl.py,sha256=1pYzjXI17FDZGPTVpmbaBvMYpB-a6SLv-84RHXA4JEA,55142
+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.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,,