| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from sentence_transformers import SentenceTransformer |
|
|
| class RouterNetwork(nn.Module): |
| def __init__(self, embedding_dim = 384, hidden_dim = 128, num_classes = 3): |
| super(RouterNetwork, self).__init__() |
| |
| |
| self.shared_layer = nn.Sequential( |
| nn.Linear(embedding_dim, hidden_dim), |
| nn.ReLU(), |
| nn.LayerNorm(hidden_dim) |
| ) |
| |
| |
| |
| self.actor_head = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim // 2), |
| nn.ReLU(), |
| nn.Linear(hidden_dim // 2, num_classes) |
| ) |
| |
| |
| self.critic_head = nn.Sequential( |
| nn.Linear(hidden_dim, hidden_dim // 2), |
| nn.ReLU(), |
| nn.Linear(hidden_dim // 2, 1) |
| ) |
| |
| def forward(self, x): |
| shared_features = self.shared_layer(x) |
| action_logits = self.actor_head(shared_features) |
| state_value = self.critic_head(shared_features) |
| return action_logits, state_value |
| |
| def get_action(self, x): |
| action_logits, value = self.forward(x) |
| action_probs = F.softmax(action_logits, dim=-1) |
| dist = torch.distributions.Categorical(action_probs) |
| action = dist.sample() |
| return action, dist.log_prob(action), value |
