-
Notifications
You must be signed in to change notification settings - Fork 0
/
dqn.py
executable file
·129 lines (115 loc) · 4.91 KB
/
dqn.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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import numpy as np
import gym
import os
import torch
import torch.nn as nn
import torch.optim as optim
import buffer
import random
import time
class DQN(nn.Module):
def __init__(self, n_input, n_output):
super(DQN, self).__init__()
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(self.device)
# Define the model
self.model = nn.Sequential(
nn.Linear(n_input, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, n_output)
).to(self.device)
# Define the target model
self.target_model = nn.Sequential(
nn.Linear(n_input, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, n_output)
).to(self.device)
self.target_update()
self.input_shape = n_input
# Define loss function and optimizer
self.loss_fn = nn.MSELoss()
self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
def forward(self, x):
return self.model(x)
def replay(self, memory, gamma, sample_size):
states, actions, rewards, dones, next_states = memory.sample(sample_size)
# Preprocessing
states = torch.from_numpy(np.asarray(states).reshape(-1, self.input_shape)).float().to(self.device)
next_states = torch.from_numpy(np.asarray(next_states).reshape(-1, self.input_shape)).float().to(self.device)
actions = torch.from_numpy(np.asarray(actions)).long().to(self.device)
rewards = torch.from_numpy(np.asarray(rewards)).float().to(self.device)
dones = torch.from_numpy(np.asarray(dones)).float().to(self.device)
# Calculate Q(s,a), shape (batch_size, action_size)
q_values = self.model(states)
# Choose only Q(s,a) of chosen actions, shape (batch_size, 1)
q_values = q_values.gather(1, actions.unsqueeze(1)).squeeze(1)
# Calculate Q(s*,a), shape (batch_size, action_size)
q_next = self.target_model(next_states)
# Choose maximum of next Q(s*,a)
q_next_max = q_next.max(1)[0]
# Calculate target Y_t
target = rewards + gamma * q_next_max * (1 - dones)
# Update the network
loss = self.loss_fn(q_values, target.detach())
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
def target_update(self):
self.target_model.load_state_dict(self.model.state_dict())
class DDQN(nn.Module):
def __init__(self, n_input, n_output):
super(DDQN, self).__init__()
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(self.device)
# Define the model
self.model = nn.Sequential(
nn.Linear(n_input, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, n_output)
).to(self.device)
# Define the target model
self.target_model = nn.Sequential(
nn.Linear(n_input, 64),
nn.ReLU(),
nn.Linear(64, 32),
nn.ReLU(),
nn.Linear(32, n_output)
).to(self.device)
self.target_update()
self.input_shape = n_input
# Define loss function and optimizer
self.loss_fn = nn.MSELoss()
self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
def forward(self, x):
return self.model(x)
def replay(self, memory, gamma, sample_size):
states, actions, rewards, dones, next_states = memory.sample(sample_size)
# Preprocessing
states = torch.from_numpy(np.asarray(states).reshape(-1, self.input_shape)).float().to(self.device)
next_states = torch.from_numpy(np.asarray(next_states).reshape(-1, self.input_shape)).float().to(self.device)
actions = torch.from_numpy(np.asarray(actions)).long().to(self.device)
rewards = torch.from_numpy(np.asarray(rewards)).float().to(self.device)
dones = torch.from_numpy(np.asarray(dones)).float().to(self.device)
# Calculate Q(s,a), shape (batch_size, action_size)
q_values = self.model(states)
# Choose only Q(s,a) of chosen actions, shape (batch_size, 1)
q_values = q_values.gather(1, actions.unsqueeze(1)).squeeze(1)
# DDQN Novelty: The main network chooses next best action
next_best_actions = self.model(next_states).max(1)[1].unsqueeze(1)
# Use target network to calculate the Q-values for that best actions
q_next = self.target_model(next_states).gather(1, next_best_actions).squeeze(1)
# Calculate target Y_t (DDQN)
target = rewards + gamma * q_next * (1 - dones)
# Update the network
loss = self.loss_fn(q_values, target.detach())
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
def target_update(self):
self.target_model.load_state_dict(self.model.state_dict())