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

rxnn/training/rl.py

@@ -10,7 +10,7 @@ from .ddp import distributed_mean
 class RlAlgorithm(ABC):
     def __init__(self):
         super(RlAlgorithm, self).__init__()
-        self.critic_loss = nn.MSELoss()
+        self.critic_loss_fn = nn.MSELoss()
 
     @abstractmethod
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
@@ -22,7 +22,7 @@ class RlAlgorithm(ABC):
         pass
 
     def critic_loss(self, values: torch.Tensor, ref_values: torch.Tensor) -> torch.Tensor:
-        return self.critic_loss(values, ref_values)
+        return self.critic_loss_fn(values, ref_values)
 
 
 class PPOConfig(TypedDict):
@@ -55,7 +55,7 @@ class PPOAlgorithm(RlAlgorithm):
         # 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)
+        return self.critic_loss_fn(values, ref_values)
 
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
                     old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
@@ -107,29 +107,31 @@ class PPOAlgorithm(RlAlgorithm):
         return policy_loss
 
     def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor,
-                     last_value: torch.Tensor, dones: torch.Tensor):
+                     last_value: torch.Tensor, dones: torch.Tensor, last_done: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         trajectory_len, batch_size = rewards.shape
         advantages = torch.zeros_like(rewards, device=rewards.device)
         last_advantage = 0
         next_value = last_value
-        next_done = torch.zeros(batch_size, device=dones.device)  # Last state is terminal
+        next_done = last_done
         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
+        for t in reversed(range(trajectory_len)):
+            if t == trajectory_len - 1:
+                # For the last step, use the provided last_value
+                delta = rewards[t] + self.gae_gamma * next_value * (1 - next_done) - values[t]
+            else:
+                # For other steps, use the next value in the trajectory
+                delta = rewards[t] + self.gae_gamma * values[t + 1] * (1 - dones[t + 1]) - values[t]
+
+            advantages[t] = delta + self.gae_gamma * self.gae_lambda * (1 - dones[t]) * last_advantage
             last_advantage = advantages[t]
-            next_value = values[t]
             next_done = dones[t]
 
         returns = advantages + values
         return advantages, returns
 
     def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor, dones: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        advantages, ref_values = self._compute_gae(rewards[:-1], values[:-1], values[-1], dones[:-1])
+        advantages, ref_values = self._compute_gae(rewards[:-1], values[:-1], values[-1], dones[:-1], dones[-1])
         if self.use_distributed_advantage_norm:
             mean_advantage = distributed_mean(advantages.mean())
             std_advantage = distributed_mean(advantages.std())

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

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

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

@@ -19,7 +19,7 @@ rxnn/training/ddp.py,sha256=VsNBjn3cY-uUj8hbsW7oKvb0_ZKnXnJ2KgObm-Mr9i4,836
 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/rl.py,sha256=DKB7oiePVMz-8Tp3jJOzxVlPZMd7HUEOJUurCJj9f68,5974
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
@@ -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.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,,
+rxnn-0.2.36.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.36.dist-info/METADATA,sha256=TZZFu-DI5dxsjYVR3qXL9UyTY_Z53Br9zlk3PtE2mlU,25960
+rxnn-0.2.36.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.36.dist-info/RECORD,,