-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathcritics.py
85 lines (73 loc) · 4.14 KB
/
critics.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
83
84
85
import torch
import torch.nn as nn
from utils import default_config
from lnets.models.layers.dense.bjorck_linear import BjorckLinear
from lnets.models.activations.group_sort import GroupSort
class Critic4mlpBjorckGS(nn.Module):
def __init__(self, in_dim, hidden_dim):
super(Critic4mlpBjorckGS, self).__init__()
self.in_dim = in_dim
self.hidden_dim = hidden_dim
# build the layers
self.layer_1 = BjorckLinear(in_features=in_dim, out_features=hidden_dim, config=default_config())
self.activation_1 = GroupSort(num_units=2)
self.layer_2 = BjorckLinear(in_features=hidden_dim, out_features=hidden_dim, config=default_config())
self.activation_2 = GroupSort(num_units=2)
self.layer_3 = BjorckLinear(in_features=hidden_dim, out_features=hidden_dim, config=default_config())
self.activation_3 = GroupSort(num_units=2)
self.layer_4 = BjorckLinear(in_features=hidden_dim, out_features=1, config=default_config())
def forward(self, batch):
out = self.activation_1(self.layer_1(batch))
out = self.activation_2(self.layer_2(out))
out = self.activation_3(self.layer_3(out))
out = self.layer_4(out)
return out
def initialize(self, critic):
assert self.hidden_dim == critic.hidden_dim
self.layer_1.weight.data = critic.layer_1.weight.data.clone()
self.layer_1.bias.data = critic.layer_1.bias.data.clone()
self.layer_2.weight.data = critic.layer_2.weight.data.clone()
self.layer_2.bias.data = critic.layer_2.bias.data.clone()
self.layer_3.weight.data = critic.layer_3.weight.data.clone()
self.layer_3.bias.data = critic.layer_3.bias.data.clone()
self.layer_4.weight.data = critic.layer_4.weight.data.clone()
self.layer_4.bias.data = critic.layer_4.bias.data.clone()
class Critic6mlpBjorckGS(nn.Module):
def __init__(self, in_dim, hidden_dim):
super(Critic6mlpBjorckGS, self).__init__()
self.in_dim = in_dim
self.hidden_dim = hidden_dim
# build the layers
self.layer_1 = BjorckLinear(in_features=in_dim, out_features=hidden_dim, config=default_config())
self.activation_1 = GroupSort(num_units=2)
self.layer_2 = BjorckLinear(in_features=hidden_dim, out_features=hidden_dim, config=default_config())
self.activation_2 = GroupSort(num_units=2)
self.layer_3 = BjorckLinear(in_features=hidden_dim, out_features=hidden_dim, config=default_config())
self.activation_3 = GroupSort(num_units=2)
self.layer_4 = BjorckLinear(in_features=hidden_dim, out_features=hidden_dim, config=default_config())
self.activation_4 = GroupSort(num_units=2)
self.layer_5 = BjorckLinear(in_features=hidden_dim, out_features=hidden_dim, config=default_config())
self.activation_5 = GroupSort(num_units=2)
self.layer_6 = BjorckLinear(in_features=hidden_dim, out_features=1, config=default_config())
def forward(self, batch):
out = self.activation_1(self.layer_1(batch))
out = self.activation_2(self.layer_2(out))
out = self.activation_3(self.layer_3(out))
out = self.activation_4(self.layer_4(out))
out = self.activation_5(self.layer_5(out))
out = self.layer_6(out)
return out
def initialize(self, critic):
assert self.hidden_dim == critic.hidden_dim
self.layer_1.weight.data = critic.layer_1.weight.data.clone()
self.layer_1.bias.data = critic.layer_1.bias.data.clone()
self.layer_2.weight.data = critic.layer_2.weight.data.clone()
self.layer_2.bias.data = critic.layer_2.bias.data.clone()
self.layer_3.weight.data = critic.layer_3.weight.data.clone()
self.layer_3.bias.data = critic.layer_3.bias.data.clone()
self.layer_4.weight.data = critic.layer_4.weight.data.clone()
self.layer_4.bias.data = critic.layer_4.bias.data.clone()
self.layer_5.weight.data = critic.layer_5.weight.data.clone()
self.layer_5.bias.data = critic.layer_5.bias.data.clone()
self.layer_6.weight.data = critic.layer_6.weight.data.clone()
self.layer_6.bias.data = critic.layer_6.bias.data.clone()