rxnn 0.2.35__tar.gz → 0.2.36__tar.gz

{rxnn-0.2.35 → rxnn-0.2.36}/PKG-INFO

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

{rxnn-0.2.35 → rxnn-0.2.36}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.35"
+version = "0.2.36"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.35 → rxnn-0.2.36}/src/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
         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())