mini CNN structure input -> (conv2d + relu) x 3 -> flatten -> (dense + relu) x 2 -> output
| 411 | |
| 412 | |
| 413 | class MarioNet(nn.Module): |
| 414 | """mini CNN structure |
| 415 | input -> (conv2d + relu) x 3 -> flatten -> (dense + relu) x 2 -> output |
| 416 | """ |
| 417 | |
| 418 | def __init__(self, input_dim, output_dim): |
| 419 | super().__init__() |
| 420 | c, h, w = input_dim |
| 421 | |
| 422 | if h != 84: |
| 423 | raise ValueError(f"Expecting input height: 84, got: {h}") |
| 424 | if w != 84: |
| 425 | raise ValueError(f"Expecting input width: 84, got: {w}") |
| 426 | |
| 427 | self.online = self.__build_cnn(c, output_dim) |
| 428 | |
| 429 | self.target = self.__build_cnn(c, output_dim) |
| 430 | self.target.load_state_dict(self.online.state_dict()) |
| 431 | |
| 432 | # Q_target parameters are frozen. |
| 433 | for p in self.target.parameters(): |
| 434 | p.requires_grad = False |
| 435 | |
| 436 | def forward(self, input, model): |
| 437 | if model == "online": |
| 438 | return self.online(input) |
| 439 | elif model == "target": |
| 440 | return self.target(input) |
| 441 | |
| 442 | def __build_cnn(self, c, output_dim): |
| 443 | return nn.Sequential( |
| 444 | nn.Conv2d(in_channels=c, out_channels=32, kernel_size=8, stride=4), |
| 445 | nn.ReLU(), |
| 446 | nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2), |
| 447 | nn.ReLU(), |
| 448 | nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1), |
| 449 | nn.ReLU(), |
| 450 | nn.Flatten(), |
| 451 | nn.Linear(3136, 512), |
| 452 | nn.ReLU(), |
| 453 | nn.Linear(512, output_dim), |
| 454 | ) |
| 455 | |
| 456 | |
| 457 | ###################################################################### |