File size: 1,557 Bytes
975dbdd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | 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__()
# shared representation layer
self.shared_layer = nn.Sequential(
nn.Linear(embedding_dim, hidden_dim),
nn.ReLU(),
nn.LayerNorm(hidden_dim)
)
# Actor Head: outputs logits for the 3 modules
# 0: unstructured, 1: website, 2: structured
self.actor_head = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim // 2),
nn.ReLU(),
nn.Linear(hidden_dim // 2, num_classes)
)
# Critic Head: outputs a single value for the state value estimation
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 |