class DQNAgent:
def __init__(
self,
env,
memory_size,
batch_size,
target_update=100,
gamma=0.99,
# replay parameters
alpha=0.2,
beta=0.6,
prior_eps=1e-6,
# Categorical DQN parameters
v_min=0,
v_max=200,
atom_size=51,
# N-step Learning
n_step=3,
start_train=32,
save_weights=True,
log=True,
lr=0.001,
seed=0,
episodes=200):
self.env = env
obs_dim = self.env.observation_dim
action_dim = self.env.action_dim
self.batch_size = batch_size
self.target_update = target_update
self.gamma = gamma
self.lr = lr
self.memory_size = memory_size
self.seed = seed
# device: cpu / gpu
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
print(self.device)
# memory for 1-step Learning
self.beta = beta
self.prior_eps = prior_eps
self.memory = PrioritizedReplayBuffer(obs_dim,
memory_size,
batch_size,
alpha=alpha)
# memory for N-step Learning
self.use_n_step = True if n_step > 1 else False
if self.use_n_step:
self.n_step = n_step
self.memory_n = ReplayBuffer(obs_dim,
memory_size,
batch_size,
n_step=n_step,
gamma=gamma)
# Categorical DQN parameters
self.v_min = v_min
self.v_max = v_max
self.atom_size = atom_size
self.support = torch.linspace(self.v_min, self.v_max,
self.atom_size).to(self.device)
# networks: dqn, dqn_target
self.dqn = Network(obs_dim, action_dim, self.atom_size,
self.support).to(self.device)
self.dqn_target = Network(obs_dim, action_dim, self.atom_size,
self.support).to(self.device)
self.dqn_target.load_state_dict(self.dqn.state_dict())
self.dqn_target.eval()
# optimizer
self.optimizer = optim.Adam(self.dqn.parameters(), lr=self.lr)
# transition to store in memory
self.transition = list()
self.fig, (self.ax1, self.ax2) = plt.subplots(2, figsize=(10, 10))
self.start_train = start_train
self.save_weights = save_weights
self.time = datetime.datetime.now().timetuple()
self.path = f"weights/{self.time[2]}-{self.time[1]}-{self.time[0]}_{self.time[3]}-{self.time[4]}"
self.log = log
self.episode_cnt = 0
self.episodes = episodes
if self.save_weights is True:
self.create_save_directory()
plt.ion()
def create_save_directory(self):
try:
os.mkdir(self.path)
except OSError:
print("Creation of the directory %s failed" % self.path)
else:
print("Successfully created the directory %s " % self.path)
def select_action(self, state):
"""Select an action from the input state."""
# NoisyNet: no epsilon greedy action selection
selected_action = self.dqn(torch.FloatTensor(state).to(
self.device)).argmax()
selected_action = selected_action.detach().cpu().numpy()
self.transition = [state, selected_action]
return selected_action
def step(self, action):
"""Take an action and return the response of the env."""
next_state, reward, done = self.env.step(action)
self.transition += [reward, next_state, done]
# N-step transition
if self.use_n_step:
one_step_transition = self.memory_n.store(*self.transition)
# 1-step transition
else:
one_step_transition = self.transition
# add a single step transition
if one_step_transition:
self.memory.store(*one_step_transition)
return next_state, reward, done
def update_model(self):
"""Update the model by gradient descent."""
# PER needs beta to calculate weights
samples = self.memory.sample_batch(self.beta)
weights = torch.FloatTensor(samples["weights"].reshape(-1, 1)).to(
self.device)
indices = samples["indices"]
# 1-step Learning loss
elementwise_loss = self._compute_dqn_loss(samples, self.gamma)
# PER: importance sampling before average
loss = torch.mean(elementwise_loss * weights)
# N-step Learning loss
# we are gonna combine 1-step loss and n-step loss so as to
# prevent high-variance. The original rainbow employs n-step loss only.
if self.use_n_step:
gamma = self.gamma**self.n_step
samples = self.memory_n.sample_batch_from_idxs(indices)
elementwise_loss_n_loss = self._compute_dqn_loss(samples, gamma)
elementwise_loss += elementwise_loss_n_loss
# PER: importance sampling before average
loss = torch.mean(elementwise_loss * weights)
self.optimizer.zero_grad()
loss.backward()
# print(loss)
clip_grad_norm_(self.dqn.parameters(), 10.0)
self.optimizer.step()
# PER: update priorities
loss_for_prior = elementwise_loss.detach().cpu().numpy()
new_priorities = loss_for_prior + self.prior_eps
self.memory.update_priorities(indices, new_priorities)
# NoisyNet: reset noise
self.dqn.reset_noise()
self.dqn_target.reset_noise()
return loss.item()
def train(self, num_frames, plotting_interval=100):
"""Train the agent."""
if self.log:
pass
# config = {'gamma': self.gamma, 'log_interval': plotting_interval, 'learning_rate': self.lr,
# 'directory': self.path, 'type': 'dqn', 'replay_memory': self.memory_size, 'environment': 'normal', 'seed': self.seed}
# wandb.init(project='is_os', entity='pydqn', config=config, notes=self.env.reward_function, reinit=True, tags=['report'])
# wandb.watch(self.dqn)
self.env.reset()
state = self.env.get_state()
won = False
update_cnt = 0
losses = []
scores = []
score = 0
frame_cnt = 0
self.episode_cnt = 0
for frame_idx in range(1, num_frames + 1):
frame_cnt += 1
action = self.select_action(state)
next_state, reward, done = self.step(action)
state = next_state
score += reward
fraction = min(frame_cnt / num_frames, 1.0)
self.beta = self.beta + fraction * (1.0 - self.beta)
# if agent has trained 500 frames, terminate
if frame_cnt == 500:
done = True
# if episode ends
if done:
if reward > 0:
won = True
self.env.reset()
state = self.env.get_state()
self.episode_cnt += 1
scores.append(score)
score = 0
frame_cnt = 0
# if training is ready
if len(self.memory) >= self.batch_size:
loss = self.update_model()
losses.append(loss)
update_cnt += 1
# if hard update is needed
if update_cnt % self.target_update == 0:
self._target_hard_update()
# plotting
if frame_idx % plotting_interval == 0:
self._plot(frame_idx, scores, losses)
if frame_idx % 1000 == 0:
torch.save(self.dqn.state_dict(),
f'{self.path}/{frame_idx}.tar')
print(f"model saved at:\n {self.path}/{frame_idx}.tar")
# wandb.run.summary['won'] = won
self.env.close()
def _compute_dqn_loss(self, samples, gamma):
"""Return categorical dqn loss."""
device = self.device # for shortening the following lines
state = torch.FloatTensor(samples["obs"]).to(device)
next_state = torch.FloatTensor(samples["next_obs"]).to(device)
action = torch.LongTensor(samples["acts"]).to(device)
reward = torch.FloatTensor(samples["rews"].reshape(-1, 1)).to(device)
done = torch.FloatTensor(samples["done"].reshape(-1, 1)).to(device)
# Categorical DQN algorithm
delta_z = float(self.v_max - self.v_min) / (self.atom_size - 1)
with torch.no_grad():
# Double DQN
next_action = self.dqn(next_state).argmax(1)
next_dist = self.dqn_target.dist(next_state)
next_dist = next_dist[range(self.batch_size), next_action]
t_z = reward + (1 - done) * gamma * self.support
t_z = t_z.clamp(min=self.v_min, max=self.v_max)
b = (t_z - self.v_min) / delta_z
l = b.floor().long()
u = b.ceil().long()
offset = (torch.linspace(
0, (self.batch_size - 1) * self.atom_size,
self.batch_size).long().unsqueeze(1).expand(
self.batch_size, self.atom_size).to(self.device))
proj_dist = torch.zeros(next_dist.size(), device=self.device)
proj_dist.view(-1).index_add_(0, (l + offset).view(-1),
(next_dist *
(u.float() - b)).view(-1))
proj_dist.view(-1).index_add_(0, (u + offset).view(-1),
(next_dist *
(b - l.float())).view(-1))
dist = self.dqn.dist(state)
log_p = torch.log(dist[range(self.batch_size), action])
elementwise_loss = -(proj_dist * log_p).sum(1)
return elementwise_loss
def _target_hard_update(self):
"""Hard update: target <- local."""
self.dqn_target.load_state_dict(self.dqn.state_dict())
def _plot(self, frame_cnt, scores, losses):
self.ax1.cla()
self.ax1.set_title(
f'frames: {frame_cnt} score: {np.mean(scores[-10:])}')
self.ax1.plot(scores[-999:], color='red')
self.ax2.cla()
self.ax2.set_title(f'loss: {np.mean(losses[-10:])}')
self.ax2.plot(losses[-999:], color='blue')
plt.show()
plt.pause(0.1)
# needed for wandb to not log nans
# if frame_cnt < self.start_train + 11:
# loss = 0
# else:
# loss = np.mean(losses[-10:])
if self.log:
pass
class Agent:
def __init__(
self,
env: 'Environment',
input_frame: ('int: the number of channels of input image'),
input_dim: (
'int: the width and height of pre-processed input image'),
input_type: ('str: the type of input dimension'),
num_frames: ('int: Total number of frames'),
skipped_frame: ('int: The number of skipped frames'),
eps_decay: ('float: Epsilon Decay_rate'),
gamma: ('float: Discount Factor'),
target_update_freq: ('int: Target Update Frequency (by frames)'),
update_type: (
'str: Update type for target network. Hard or Soft') = 'hard',
soft_update_tau: ('float: Soft update ratio') = None,
batch_size: ('int: Update batch size') = 32,
buffer_size: ('int: Replay buffer size') = 1000000,
alpha: (
'int: Hyperparameter for how large prioritization is applied'
) = 0.5,
beta:
('int: Hyperparameter for the annealing factor of importance sampling'
) = 0.5,
epsilon_for_priority:
('float: Hyperparameter for adding small increment to the priority'
) = 1e-6,
update_start_buffer_size: (
'int: Update starting buffer size') = 50000,
learning_rate: ('float: Learning rate') = 0.0004,
eps_min: ('float: Epsilon Min') = 0.1,
eps_max: ('float: Epsilon Max') = 1.0,
device_num: ('int: GPU device number') = 0,
rand_seed: ('int: Random seed') = None,
plot_option: ('str: Plotting option') = False,
model_path: ('str: Model saving path') = './'):
self.action_dim = env.action_space.n
self.device = torch.device(
f'cuda:{device_num}' if torch.cuda.is_available() else 'cpu')
self.model_path = model_path
self.env = env
self.input_frames = input_frame
self.input_dim = input_dim
self.num_frames = num_frames
self.skipped_frame = skipped_frame
self.epsilon = eps_max
self.eps_decay = eps_decay
self.eps_min = eps_min
self.gamma = gamma
self.target_update_freq = target_update_freq
self.update_cnt = 0
self.update_type = update_type
self.tau = soft_update_tau
self.batch_size = batch_size
self.buffer_size = buffer_size
self.update_start = update_start_buffer_size
self.seed = rand_seed
self.plot_option = plot_option
# hyper parameters for PER
self.alpha = alpha
self.beta = beta
self.beta_step = (1.0 - beta) / num_frames
self.epsilon_for_priority = epsilon_for_priority
if input_type == '1-dim':
self.q_current = QNetwork_1dim(self.input_dim,
self.action_dim).to(self.device)
self.q_target = QNetwork_1dim(self.input_dim,
self.action_dim).to(self.device)
else:
self.q_current = QNetwork(
(self.input_frames, self.input_dim, self.input_dim),
self.action_dim).to(self.device)
self.q_target = QNetwork(
(self.input_frames, self.input_dim, self.input_dim),
self.action_dim).to(self.device)
self.q_target.load_state_dict(self.q_current.state_dict())
self.q_target.eval()
self.optimizer = optim.Adam(self.q_current.parameters(),
lr=learning_rate)
if input_type == '1-dim':
self.memory = PrioritizedReplayBuffer(self.buffer_size,
self.input_dim,
self.batch_size, self.alpha,
input_type)
else:
self.memory = PrioritizedReplayBuffer(
self.buffer_size,
(self.input_frames, self.input_dim, self.input_dim),
self.batch_size, self.alpha, input_type)
def select_action(
self, state:
'Must be pre-processed in the same way while updating current Q network. See def _compute_loss'
):
if np.random.random() < self.epsilon:
return np.zeros(self.action_dim), self.env.action_space.sample()
else:
state = torch.FloatTensor(state).to(self.device).unsqueeze(0)
Qs = self.q_current(state)
action = Qs.argmax()
return Qs.detach().cpu().numpy(), action.detach().item()
def processing_resize_and_gray(self, frame):
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) # Pure
# frame = cv2.cvtColor(frame[:177, 32:128, :], cv2.COLOR_RGB2GRAY) # Boxing
# frame = cv2.cvtColor(frame[2:198, 7:-7, :], cv2.COLOR_RGB2GRAY) # Breakout
frame = cv2.resize(frame,
dsize=(self.input_dim, self.input_dim)).reshape(
self.input_dim, self.input_dim).astype(np.uint8)
return frame
def get_state(self, state, action, skipped_frame=0):
'''
num_frames: how many frames to be merged
input_size: hight and width of input resized image
skipped_frame: how many frames to be skipped
'''
next_state = np.zeros(
(self.input_frames, self.input_dim, self.input_dim))
for i in range(len(state) - 1):
next_state[i] = state[i + 1]
rewards = 0
dones = 0
for j in range(skipped_frame):
state, reward, done, _ = self.env.step(action)
rewards += reward
dones += int(done)
state, reward, done, _ = self.env.step(action)
next_state[-1] = self.processing_resize_and_gray(state)
rewards += reward
dones += int(done)
return rewards, next_state, dones
def get_state_1dim(self, state, action, skipped_frame=0):
'''
num_frames: how many frames to be merged
input_size: hight and width of input resized image
skipped_frame: how many frames to be skipped
'''
next_state = np.zeros((self.input_frames, self.input_dim))
for i in range(len(state) - 1):
next_state[i] = state[i + 1]
rewards = 0
dones = 0
for _ in range(skipped_frame):
state, reward, done, _ = self.env.step(action)
rewards += reward
dones += int(done)
state, reward, done, _ = self.env.step(action)
next_state[-1] = state
rewards += reward
dones += int(done)
return rewards, next_state, dones
def get_init_state(self):
init_state = np.zeros(
(self.input_frames, self.input_dim, self.input_dim))
init_frame = self.env.reset()
init_state[0] = self.processing_resize_and_gray(init_frame)
for i in range(1, self.input_frames):
action = self.env.action_space.sample()
for j in range(self.skipped_frame):
state, _, _, _ = self.env.step(action)
state, _, _, _ = self.env.step(action)
init_state[i] = self.processing_resize_and_gray(state)
return init_state
def get_init_state_1dim(self):
init_state = np.zeros((self.input_frames, self.input_dim))
init_frame = self.env.reset()
init_state[0] = init_frame
for i in range(1, self.input_frames):
action = self.env.action_space.sample()
for j in range(self.skipped_frame):
state, _, _, _ = self.env.step(action)
state, _, _, _ = self.env.step(action)
init_state[i] = state
return init_state
def store(self, state, action, reward, next_state, done):
self.memory.store(state, action, reward, next_state, done)
def update_current_q_net(self):
'''The diffent method between Dueling and PER in the Agent class'''
batch = self.memory.batch_load(self.beta)
weights = torch.FloatTensor(batch['weights'].reshape(-1, 1)).to(
self.device)
sample_wise_loss = self._compute_loss(
batch) # PER: shape of loss -> (batch, 1)
loss = torch.mean(sample_wise_loss * weights)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# For PER: update priorities of the samples.
sample_wise_loss = sample_wise_loss.detach().cpu().numpy()
batch_priorities = sample_wise_loss + self.epsilon_for_priority
self.memory.update_priorities(batch['indices'], batch_priorities)
return loss.item()
def target_soft_update(self):
for target_param, current_param in zip(self.q_target.parameters(),
self.q_current.parameters()):
target_param.data.copy_(self.tau * current_param.data +
(1.0 - self.tau) * target_param.data)
def target_hard_update(self):
self.update_cnt = (self.update_cnt + 1) % self.target_update_freq
if self.update_cnt == 0:
self.q_target.load_state_dict(self.q_current.state_dict())
def train(self):
tic = time.time()
losses = []
scores = []
epsilons = []
avg_scores = [[-1000]]
score = 0
print("Storing initial buffer..")
# state = self.get_init_state()
# state = self.get_init_state_1dim()
state = self.env.reset()
for frame_idx in range(1, self.update_start + 1):
_, action = self.select_action(state)
next_state, reward, done, _ = self.env.step(action)
self.store(state, action, reward, next_state, done)
state = next_state
if done: state = self.env.reset()
print("Done. Start learning..")
history_store = []
for frame_idx in range(1, self.num_frames + 1):
Qs, action = self.select_action(state)
next_state, reward, done, _ = self.env.step(action)
self.store(state, action, reward, next_state, done)
history_store.append([state, Qs, action, reward, next_state, done])
loss = self.update_current_q_net()
if self.update_type == 'hard': self.target_hard_update()
elif self.update_type == 'soft': self.target_soft_update()
score += reward
losses.append(loss)
if done:
scores.append(score)
if np.mean(scores[-10:]) > max(avg_scores):
torch.save(
self.q_current.state_dict(),
self.model_path + '{}_Score:{}.pt'.format(
frame_idx, np.mean(scores[-10:])))
training_time = round((time.time() - tic) / 3600, 1)
np.save(
self.model_path +
'{}_history_Score_{}_{}hrs.npy'.format(
frame_idx, score, training_time),
np.array(history_store))
print(
" | Model saved. Recent scores: {}, Training time: {}hrs"
.format(scores[-10:], training_time),
' /'.join(os.getcwd().split('/')[-3:]))
avg_scores.append(np.mean(scores[-10:]))
if self.plot_option == 'inline':
scores.append(score)
epsilons.append(self.epsilon)
self._plot(frame_idx, scores, losses, epsilons)
elif self.plot_option == 'wandb':
Q_mean = np.mean(np.array(history_store)[:, 1])
wandb.log({
'Score': score,
'loss(10 frames avg)': np.mean(losses[-10:]),
'Q (mean)': Q_mean,
'Epsilon': self.epsilon,
'beta': self.beta
})
print(score, end='\r')
else:
print(score, end='\r')
score = 0
state = self.env.reset()
history_store = []
else:
state = next_state
self._epsilon_step()
# self.beta = min(self.beta+self.beta_step, 1.0) # for PER. beta is increased linearly up to 1.0
fraction = min(frame_idx / self.num_frames, 1.0)
self.beta = self.beta + fraction * (1.0 - self.beta)
print("Total training time: {}(hrs)".format(
(time.time() - tic) / 3600))
def _epsilon_step(self):
''' Epsilon decay control '''
eps_decay_init = 1 / 1200000
eps_decay = [
eps_decay_init, eps_decay_init / 2.5, eps_decay_init / 3.5,
eps_decay_init / 5.5
]
if self.epsilon > 0.30:
self.epsilon = max(self.epsilon - eps_decay[0], 0.1)
elif self.epsilon > 0.27:
self.epsilon = max(self.epsilon - eps_decay[1], 0.1)
elif self.epsilon > 1.7:
self.epsilon = max(self.epsilon - eps_decay[2], 0.1)
else:
self.epsilon = max(self.epsilon - eps_decay[3], 0.1)
def _compute_loss(self, batch: "Dictionary (S, A, R', S', Dones)"):
# If normalization is used, it must be applied to 'state' and 'next_state' here. ex) state/255
states = torch.FloatTensor(batch['states']).to(self.device)
next_states = torch.FloatTensor(batch['next_states']).to(self.device)
actions = torch.LongTensor(batch['actions'].reshape(-1,
1)).to(self.device)
rewards = torch.FloatTensor(batch['rewards'].reshape(-1, 1)).to(
self.device)
dones = torch.FloatTensor(batch['dones'].reshape(-1,
1)).to(self.device)
current_q = self.q_current(states).gather(1, actions)
# The next line is the only difference from Vanila DQN.
next_q = self.q_target(next_states).gather(
1,
self.q_current(next_states).argmax(axis=1, keepdim=True)).detach()
mask = 1 - dones
target = (rewards + (mask * self.gamma * next_q)).to(self.device)
# For PER, the shape of loss is (batch, 1). Therefore, using "reduction='none'" option.
sample_wise_loss = F.smooth_l1_loss(current_q,
target,
reduction="none")
return sample_wise_loss
def _plot(self, frame_idx, scores, losses, epsilons):
clear_output(True)
plt.figure(figsize=(20, 5), facecolor='w')
plt.subplot(131)
plt.title('frame %s. score: %s' % (frame_idx, np.mean(scores[-10:])))
plt.plot(scores)
plt.subplot(132)
plt.title('loss')
plt.plot(losses)
plt.subplot(133)
plt.title('epsilons')
plt.plot(epsilons)
plt.show()