rxnn 0.2.12__tar.gz → 0.2.13__tar.gz

{rxnn-0.2.12 → rxnn-0.2.13}/PKG-INFO

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

{rxnn-0.2.12 → rxnn-0.2.13}/pyproject.toml

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

{rxnn-0.2.12 → rxnn-0.2.13}/src/rxnn/memory/stm.py

@@ -39,6 +39,7 @@ class ShortTermMemory(nn.Module):
         return self.memory[layer]
 
     def update_layer(self, layer: int, new_stm: torch.Tensor):
+        self.memory = self.memory.clone()
         self.memory[layer] = new_stm
 
     def update_all(self, new_stm: torch.Tensor):
@@ -60,7 +61,7 @@ class ShortTermMemory(nn.Module):
             self.register_buffer('memory', trained_stm)
 
     def reset(self, init_type: str = None):
-        self.memory.copy_(self._init_tensor(init_type))
+        self.memory = self._init_tensor(init_type).to(self.memory.device)
 
     def resize(self, new_stm_size: int, init_type: str = None):
         self.stm_size = new_stm_size
@@ -85,8 +86,7 @@ class ShortTermMemory(nn.Module):
         if use_mean_from_batch:
             batch_mean = self.memory.mean(dim=(1, 2, 3), keepdim=True)
             delattr(self, 'memory')
-            self.register_buffer('memory', self._init_tensor())
-            self.memory.copy_(batch_mean)
+            self.register_buffer('memory', batch_mean)
         else:
             delattr(self, 'memory')
             self.register_buffer('memory', self._init_tensor())

{rxnn-0.2.12 → rxnn-0.2.13}/src/rxnn/training/mrl.py

@@ -461,7 +461,6 @@ class MRLTrainer:
         # 1. Run update separately for episodes in trajectory - we have to reset memory before each episode, and update
         # memory, based on collected episode data
         all_losses = []
-        trajectories_len = len(trajectories)
         for episode_idx, episode in enumerate(trajectories):
             episode_steps = episode['steps']
             should_reset_stm = episode['reset_stm']
@@ -514,14 +513,14 @@ class MRLTrainer:
 
                 # 9. Update the model in AMP or regular mode
                 if self.use_amp:
-                    self.scaler.scale(policy_loss).backward()
+                    self.scaler.scale(policy_loss).backward(retain_graph=True)
                     self.scaler.unscale_(self.optimizer)
                     torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                                    error_if_nonfinite=False)
                     self.scaler.step(self.optimizer)
                     self.scaler.update()
                 else:
-                    policy_loss.backward()
+                    policy_loss.backward(retain_graph=True)
                     torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                                    error_if_nonfinite=False)
                     self.optimizer.step()

{rxnn-0.2.12 → rxnn-0.2.13}/src/rxnn/training/rl.py

@@ -43,6 +43,8 @@ class PPOAlgorithm(RlAlgorithm):
         # b) Calculate ratio
         ratio = (new_log_probs - old_log_probs).exp()
 
+        advantages = advantages.unsqueeze(-1)
+
         # c) Clipped surrogate loss
         surr1 = ratio * advantages
         surr2 = torch.clamp(ratio, 1.0 - self.clip_eps, 1.0 + self.clip_eps) * advantages