rxnn 0.2.35__py3-none-any.whl → 0.2.37__py3-none-any.whl

rxnn/training/rl.py

@@ -107,12 +107,12 @@ 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.float()
         dones = dones.float()
 
         for t in reversed(range(trajectory_len)):
@@ -131,7 +131,7 @@ class PPOAlgorithm(RlAlgorithm):
         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.35.dist-info → rxnn-0.2.37.dist-info}/METADATA

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

{rxnn-0.2.35.dist-info → rxnn-0.2.37.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=OyfqVEh5Z-YvoW7baTTpOH0DuYexUOfhkZmbCpEoS_A,5962
+rxnn/training/rl.py,sha256=47wxFeUSHSqc1dKEEy8skTcNHDqNuthsYTGA-HeUbhg,5982
 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.35.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.35.dist-info/METADATA,sha256=EpcmlYIdyJMIP6KnT-065MWoGUudGYFoSbeRJiZAqnk,25960
-rxnn-0.2.35.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.35.dist-info/RECORD,,
+rxnn-0.2.37.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.37.dist-info/METADATA,sha256=GXBCyK-3ALJw6TpVk7rJ7Z2uyFSq8u8N-TYpmQaeUE8,25960
+rxnn-0.2.37.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.37.dist-info/RECORD,,