(self, batch=True)
| 44 | See https://github.com/pytorch/examples/tree/main/reinforcement_learning |
| 45 | """ |
| 46 | def __init__(self, batch=True): |
| 47 | super(Policy, self).__init__() |
| 48 | self.affine1 = nn.Linear(4, 128) |
| 49 | self.dropout = nn.Dropout(p=0.6) |
| 50 | self.affine2 = nn.Linear(128, 2) |
| 51 | self.dim = 2 if batch else 1 |
| 52 | |
| 53 | def forward(self, x): |
| 54 | x = self.affine1(x) |