rxnn 0.2.20__tar.gz → 0.2.22__tar.gz

{rxnn-0.2.20 → rxnn-0.2.22}/PKG-INFO

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

{rxnn-0.2.20 → rxnn-0.2.22}/pyproject.toml

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

{rxnn-0.2.20 → rxnn-0.2.22}/src/rxnn/training/models.py

@@ -129,7 +129,7 @@ class MrlActorModel(nn.Module):
             _, ed = self.encoder(x, attention_mask=attention_mask)
             return self.memory_attention(ed, attention_mask=attention_mask)
 
-class MrlCriticModel(nn.Module):
+class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipeline_tag="text-classification"):
     def __init__(self, encoder: nn.Module, embed_dim: int, **kwargs):
         super(MrlCriticModel, self).__init__(**kwargs)
         self.encoder = encoder

{rxnn-0.2.20 → rxnn-0.2.22}/src/rxnn/training/rl.py

@@ -24,8 +24,6 @@ class RlAlgorithm(ABC):
         return self.critic_loss(rewards, values)
 
 class PPOConfig(TypedDict):
-    gae_gamma: float
-    gae_lambda: float
     clip_eps: float
 
 class PPOAlgorithm(RlAlgorithm):
@@ -33,8 +31,6 @@ class PPOAlgorithm(RlAlgorithm):
         super(PPOAlgorithm, self).__init__()
 
         # PPO Config
-        self.gae_gamma = config.get('gae_gamma', 0.99)
-        self.gae_lambda = config.get('gae_lambda', 0.95)
         self.clip_eps = config.get('clip_eps', 0.2)
 
     def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
@@ -86,15 +82,6 @@ class PPOAlgorithm(RlAlgorithm):
 
         return policy_loss
 
-    # def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, next_value: torch.Tensor) -> torch.Tensor:
-    #     advantages = torch.zeros_like(rewards, device=values.device)
-    #     last_advantage = 0
-    #     for t in reversed(range(rewards.size(0))):
-    #         delta = rewards[t] + self.gae_gamma * next_value - values[t]
-    #         advantages[t] = delta + self.gae_gamma * self.gae_lambda * last_advantage
-    #         last_advantage = advantages[t]
-    #     return advantages
-
     def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
         advantages = rewards - values
         normalized_advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)