| 29 | |
| 30 | |
| 31 | class Policy(nn.Module): |
| 32 | def __init__(self): |
| 33 | super(Policy, self).__init__() |
| 34 | self.affine1 = nn.Linear(4, 128) |
| 35 | self.dropout = nn.Dropout(p=0.6) |
| 36 | self.affine2 = nn.Linear(128, 2) |
| 37 | |
| 38 | self.saved_log_probs = [] |
| 39 | self.rewards = [] |
| 40 | |
| 41 | def forward(self, x): |
| 42 | x = self.affine1(x) |
| 43 | x = self.dropout(x) |
| 44 | x = F.relu(x) |
| 45 | action_scores = self.affine2(x) |
| 46 | return F.softmax(action_scores, dim=1) |
| 47 | |
| 48 | |
| 49 | policy = Policy() |