| 33 | |
| 34 | # Policy network: outputs logits over actions. |
| 35 | class Actor(nn.Module): |
| 36 | def __init__(self, state_size, action_size): |
| 37 | super().__init__() |
| 38 | self.fc1 = nn.Linear(state_size, 24) |
| 39 | self.fc2 = nn.Linear(24, action_size) |
| 40 | nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity="relu") |
| 41 | nn.init.kaiming_uniform_(self.fc2.weight, nonlinearity="relu") |
| 42 | |
| 43 | def forward(self, x): |
| 44 | return self.fc2(torch.relu(self.fc1(x))) |
| 45 | |
| 46 | |
| 47 | # Value network: outputs a scalar V(s). |