| from __future__ import annotations | |
| import torch | |
| from schema import ACTIONS | |
| from torch import nn | |
| class NeuralModelMachine(nn.Module): | |
| def __init__(self) -> None: | |
| super().__init__() | |
| self.action_embedding = nn.Embedding(len(ACTIONS), 16) | |
| self.network = nn.Sequential( | |
| nn.Linear(8 + 4 + 16, 64), | |
| nn.LayerNorm(64), | |
| nn.SiLU(), | |
| nn.Linear(64, 64), | |
| nn.SiLU(), | |
| nn.Linear(64, 9), | |
| ) | |
| def forward( | |
| self, | |
| state: torch.Tensor, | |
| capabilities: torch.Tensor, | |
| action: torch.Tensor, | |
| ) -> torch.Tensor: | |
| action_features = self.action_embedding(action) | |
| return self.network(torch.cat([state, capabilities, action_features], dim=1)) | |
| def transition( | |
| self, | |
| state: torch.Tensor, | |
| capabilities: torch.Tensor, | |
| action: torch.Tensor, | |
| ) -> tuple[torch.Tensor, torch.Tensor]: | |
| probabilities = torch.sigmoid(self(state, capabilities, action)) | |
| return (probabilities[:, :8] >= 0.5).float(), probabilities[:, 8] | |
| def parameter_count(model: nn.Module) -> int: | |
| return sum(parameter.numel() for parameter in model.parameters()) | |