-
Notifications
You must be signed in to change notification settings - Fork 0
/
dqn_torch.py
73 lines (60 loc) · 3 KB
/
dqn_torch.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
import os, random
from copy import deepcopy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from buffer import ReplayBuffer
from network import Net
class Agent:
def __init__(self, lr=0.003, input_dims=[4], env=None, gamma=0.99, n_actions=2, epsilon_greedy_start=0.5,
epsilon_greedy_decay=0.0002, max_size=1000000, layer1_size=64, layer2_size=64, batch_size=128, writer=None):
self.env = env
self.gamma = gamma
self.memory = ReplayBuffer(max_size, input_dims, n_actions)
self.batch_size = batch_size
self.n_actions = n_actions
self.epsilon_greedy_start = epsilon_greedy_start
self.epsilon_greedy_decay = epsilon_greedy_decay
self.net = Net(lr, input_dims, n_actions=n_actions, fc1_dims=layer1_size, fc2_dims=layer2_size, name='dqn')
self.target_net = deepcopy(self.net)
self.target_net.load_state_dict(self.net.state_dict())
self.target_net.eval()
self.criterion = F.smooth_l1_loss
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.net.to(self.device)
self.target_net.to(self.device)
self.writer = writer
def choose_action(self, state, timestep):
epsilon = self.epsilon_greedy_start - self.epsilon_greedy_decay * timestep
if random.random() <= epsilon:
return self.env.action_space.sample()
state = torch.from_numpy(state).to(self.device, torch.float)
action = self.net.forward(state).max(0)[1].item()
return action
def target_update(self):
self.target_net.load_state_dict(self.net.state_dict())
def model_update(self, timestep):
if len(self.memory) < self.batch_size:
return
states, actions, rewards, states_, terminals = self.memory.sample_buffer(self.batch_size)
states = states.to(self.device)
actions = actions.to(self.device)
rewards = rewards.to(self.device)
states_ = states_.to(self.device)
terminals = terminals.to(self.device)
state_actinon_values = self.net.forward(states.to(torch.float))
state_actinon_values = state_actinon_values.gather(1, actions[:, 0].unsqueeze(1).to(torch.long)).squeeze(1)
with torch.no_grad():
next_state_values = self.target_net(states_.to(torch.float)).max(1)[0].detach()
expected_action_values = self.gamma * next_state_values + rewards
expected_action_values = expected_action_values * (1 - terminals.to(torch.uint8))
loss = self.criterion(state_actinon_values, expected_action_values.to(torch.float))
self.writer.add_scalar("loss", loss.item(), timestep)
self.net.optimizer.zero_grad()
loss.backward()
for param in self.net.parameters():
param.grad.data.clamp_(-1, 1)
self.net.optimizer.step()
def store_transition(self, state, action, reward, state_, done):
self.memory.store_transtions(state, action, reward, state_, done)