-
Notifications
You must be signed in to change notification settings - Fork 1
/
rainbow_dqn.py
147 lines (105 loc) · 5.6 KB
/
rainbow_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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
#!/usr/bin/env python3
import gym
import ptan
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from tensorboardX import SummaryWriter
from lib import common
from lib.common import *
from lib.rainbow_model import RainbowDQN
def calc_loss(batch, batch_weights, net, tgt_net, gamma, device="cpu"):
states, actions, rewards, dones, next_states = common.unpack_batch(batch)
batch_size = len(batch)
states_v = torch.tensor(states).to(device)
actions_v = torch.tensor(actions).to(device)
next_states_v = torch.tensor(next_states).to(device)
batch_weights_v = torch.tensor(batch_weights).to(device)
# next state distribution
# dueling arch -- actions from main net, distr from tgt_net
# calc at once both next and cur states
distr_v, qvals_v = net.both(torch.cat((states_v, next_states_v)))
next_qvals_v = qvals_v[batch_size:]
distr_v = distr_v[:batch_size]
next_actions_v = next_qvals_v.max(1)[1]
next_distr_v = tgt_net(next_states_v)
next_best_distr_v = next_distr_v[range(batch_size), next_actions_v.data]
next_best_distr_v = tgt_net.apply_softmax(next_best_distr_v)
next_best_distr = next_best_distr_v.data.cpu().numpy()
dones = dones.astype(np.bool)
# project our distribution using Bellman update
proj_distr = common.distr_projection(next_best_distr, rewards, dones, Vmin, Vmax, N_ATOMS, gamma)
# calculate net output
state_action_values = distr_v[range(batch_size), actions_v.data]
state_log_sm_v = F.log_softmax(state_action_values, dim=1)
proj_distr_v = torch.tensor(proj_distr).to(device)
loss_v = -state_log_sm_v * proj_distr_v
loss_v = batch_weights_v * loss_v.sum(dim=1)
return loss_v.mean(), loss_v + 1e-5
def infer_buffer_weights(net, tgt_net, batches_count, buffer, params, beta, device):
print("Preparing buffer weights in {} passes".format(batches_count))
for i in range(batches_count):
if i % 1000 == 0:
print("batch {}".format(i))
batch, batch_indices, batch_weights = buffer.sample(params['batch_size'], beta)
loss_v, sample_prios_v = calc_loss(batch, batch_weights, net, tgt_net.target_model,
params['gamma'] ** REWARD_STEPS, device=device)
buffer.update_priorities(batch_indices, sample_prios_v.data.cpu().numpy())
if __name__ == "__main__":
params = common.HYPERPARAMS['pacman']
parser = argparse.ArgumentParser()
parser.add_argument("--cuda", default=False, action="store_true", help="Enable cuda")
args = parser.parse_args()
device = torch.device("cuda" if args.cuda else "cpu")
wrapper = params["env_wrapper_train"]
env = gym.make(params['env_name'])
env = wrapper(env)
writer = SummaryWriter(comment="-" + params['run_name'] + "-rainbow")
net = RainbowDQN(env.observation_space.shape, env.action_space.n).to(device)
agent = ptan.agent.DQNAgent(lambda x: net.qvals(x), ptan.actions.ArgmaxActionSelector(), device=device)
exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, gamma=params['gamma'], steps_count=REWARD_STEPS)
buffer = ptan.experience.PrioritizedReplayBuffer(exp_source, params['replay_size'], PRIO_REPLAY_ALPHA)
optimizer = optim.Adam(net.parameters(), lr=params['learning_rate'])
frame_idx = 0
first_train_iteration = True
beta = BETA_START
#If we are requested to continue training from an old checkpoint, load it
saves_filename = params['resume_from']
if saves_filename is not None:
frame_idx = int(saves_filename.split("_")[1].split(".")[0])
print("Loading network and optimizer {}".format(saves_filename))
net.load_state_dict(torch.load(params["save_dir"] + saves_filename))
optimizer.load_state_dict(torch.load(params["save_dir"] + saves_filename + ".optimizer"))
tgt_net = ptan.agent.TargetNet(net)
with common.RewardTracker(writer, params['stop_reward']) as reward_tracker:
while True:
frame_idx += 1
buffer.populate(1)
beta = min(1.0, BETA_START + frame_idx * (1.0 - BETA_START) / BETA_FRAMES)
new_rewards = exp_source.pop_total_rewards()
if new_rewards:
if reward_tracker.reward(new_rewards[0], frame_idx):
break
if len(buffer) < params['replay_initial']:
continue
if first_train_iteration:
first_train_iteration = False
infer_buffer_weights(net, tgt_net, params['replay_initial'] // params['batch_size'] * 16, buffer, params, beta, device)
optimizer.zero_grad()
batch, batch_indices, batch_weights = buffer.sample(params['batch_size'], beta)
if frame_idx % params['qvalues_estimation_interval'] == 0:
avg_qvalues = calc_avg_qval(batch, net, device=device)
writer.add_scalar("Batch qvalues", avg_qvalues, frame_idx)
loss_v, sample_prios_v = calc_loss(batch, batch_weights, net, tgt_net.target_model,
params['gamma'] ** REWARD_STEPS, device=device)
loss_v.backward()
optimizer.step()
buffer.update_priorities(batch_indices, sample_prios_v.data.cpu().numpy())
if frame_idx % params['target_net_sync'] == 0:
tgt_net.sync()
if frame_idx % params['save_interval'] == 0:
common.save_net(net, optimizer, params['save_dir'], "{}_{}.dat".format(params['run_name'], frame_idx))
common.save_net(net, optimizer, params['save_dir'], "best.dat")