/
TorchModelClasses.py
62 lines (46 loc) · 1.73 KB
/
TorchModelClasses.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
import torch
from torch import nn, optim
from torch.nn import functional as F
import numpy as np
"""
PS : {'action_space': 4, 'observation_space': (8,)}
"""
class SimpleNet(nn.Module):
def __init__(self, ALPHA, INPUT_SHAPE, OUTPUT_SHAPE):
super().__init__()
self.ALPHA = ALPHA
self.model = nn.Sequential(
nn.Linear(INPUT_SHAPE, 64),
nn.ReLU(),
nn.Linear(64, 64),
nn.ReLU(),
nn.Linear(64, OUTPUT_SHAPE),
)
self.optimizer = optim.Adam(self.parameters(), lr=self.ALPHA) # Tune this
self.loss = nn.MSELoss()
self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
self.to(self.device)
def forward(self, observation):
observation = torch.Tensor(observation).to(self.device)
actions = self.model(observation)
return actions
class HuberNet(nn.Module):
def __init__(self, ALPHA, INPUT_SHAPE, OUTPUT_SHAPE):
super().__init__()
self.ALPHA = ALPHA
self.model = nn.Sequential(
nn.Linear(INPUT_SHAPE, 64),
nn.ReLU(),
nn.Linear(64, 64),
nn.ReLU(),
nn.Linear(64, OUTPUT_SHAPE),
)
# self.optimizer = optim.RMSprop(self.parameters(), lr=self.ALPHA, momentum=0.0001) # Tune this
self.optimizer = optim.Adam(self.parameters(), lr=self.ALPHA) # Tune this
self.loss = nn.SmoothL1Loss()
self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
self.to(self.device)
def forward(self, observation):
observation = torch.Tensor(observation).to(self.device)
actions = self.model(observation)
return actions