-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
82 lines (63 loc) · 2.42 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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from torch import nn
import torch.nn.functional as F
class Model(nn.Module):
def __init__(self, n_states, n_actions):
super(Model, self).__init__()
self.n_states = n_states
self.n_actions = n_actions
self.fc1 = nn.Linear(self.n_states, 256)
self.fc2 = nn.Linear(256, 128)
self.q_values = nn.Linear(128, self.n_actions)
nn.init.kaiming_normal_(self.fc1.weight)
self.fc1.bias.data.zero_()
nn.init.kaiming_normal_(self.fc2.weight)
self.fc2.bias.data.data.zero_()
nn.init.xavier_uniform_(self.q_values.weight)
self.q_values.bias.data.zero_()
def forward(self, inputs):
x = inputs
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
return self.q_values(x)
class RNDModel(nn.Module):
def __init__(self, n_states, n_outputs):
super(RNDModel, self).__init__()
self.n_states = n_states
self.n_outputs = n_outputs
self.fc1 = nn.Linear(self.n_states, 256)
self.fc2 = nn.Linear(256, 128)
self.encoded_features = nn.Linear(128, self.n_outputs)
nn.init.kaiming_normal_(self.fc1.weight)
self.fc1.bias.data.data.zero_()
nn.init.kaiming_normal_(self.fc2.weight)
self.fc2.bias.data.data.zero_()
nn.init.xavier_uniform_(self.encoded_features.weight)
self.encoded_features.bias.data.zero_()
def forward(self, inputs):
x = inputs
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
return self.encoded_features(x)
class CRWModel(nn.Module):
def __init__(self, n_states, n_outputs):
super(CRWModel, self).__init__()
self.n_states = n_states
self.n_outputs = n_outputs
self.fc1 = nn.Linear(self.n_states, 256)
self.fc2 = nn.Linear(256, 128)
self.encoded_features = nn.Linear(128, self.n_outputs)
self.dropout=nn.Dropout(0.3)
def weight_init(self):
nn.init.kaiming_normal_(self.fc1.weight)
self.fc1.bias.data.data.zero_()
nn.init.kaiming_normal_(self.fc2.weight)
self.fc2.bias.data.data.zero_()
nn.init.xavier_uniform_(self.encoded_features.weight)
self.encoded_features.bias.data.zero_()
def forward(self, inputs):
x = inputs
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = F.relu(self.fc2(x))
x = self.dropout(x)
return self.encoded_features(x)