-
Notifications
You must be signed in to change notification settings - Fork 3
/
model.py
57 lines (50 loc) · 1.83 KB
/
model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import torch.nn as nn
import torch.nn.functional as F
import torch
class DQN(nn.Module):
def __init__(self, in_channels, n_actions):
super().__init__()
self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
self.fc1 = nn.Linear(64*7*7, 512)
self.fc2 = nn.Linear(512, n_actions)
def forward(self, x):
'''
Input shape: (bs,c,h,w) #(bs,4,84,84)
Output shape: (bs,n_actions)
'''
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = torch.flatten(x,1)
x = F.relu(self.fc1(x))
return self.fc2(x)
# Dueling DQN
class DuelDQN(nn.Module):
def __init__(self, in_channels, n_actions):
super().__init__()
self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
self.bn1 = nn.BatchNorm2d(32)
self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
self.bn2 = nn.BatchNorm2d(64)
self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
self.bn3 = nn.BatchNorm2d(64)
self.Alinear1 = nn.Linear(64*7*7,128)
self.Alinear2 = nn.Linear(128,n_actions)
self.Vlinear1 = nn.Linear(64*7*7,128)
self.Vlinear2 = nn.Linear(128,1)
def forward(self, x):
'''
Input shape: (bs,c,h,w) #(bs,4,84,84)
Output shape: (bs,n_actions)
'''
x = F.relu(self.conv1(x))
x = F.relu(self.conv2(x))
x = F.relu(self.conv3(x))
x = torch.flatten(x,1)
Ax = F.leaky_relu(self.Alinear1(x))
Ax = self.Alinear2(Ax)
Vx = F.leaky_relu(self.Vlinear1(x))
Vx = self.Vlinear2(Vx)
return Vx + (Ax - Ax.mean())