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
a = model.randomAction() else: a = model.predictAction(s) # anneal epsilon epsilon = max(0.2, epsilon - epsilon_step) # apply action, get rewards and new state s2 s2_text, r, terminal, info = env.step(a) s2 = sent2seq(s2_text, seq_len) # add current exp to buffer replay_buffer.add(s, a, r, terminal, s2) # Keep adding experience to the memory until # there are at least minibatch size samples if ((replay_buffer.size() > args.batch_size) and (step_ctr % args.rounds_per_learn == 0)): s_batch, a_batch, r_batch, t_batch, s2_batch = \ replay_buffer.sample_batch(args.batch_size) # Update the networks each given the new target values l = model.trainOnBatch(s_batch, a_batch, r_batch, t_batch, s2_batch) loss += l step_ctr = 0 s = s2 ep_reward += r cnt_invalid_actions += 1 if r == -0.1 else 0 if terminal: break ep_lens.append(j + 1) invalids.append(cnt_invalid_actions) quests_complete.append(int(r >= 1)) scores.append(ep_reward)