def train(game, num_steps=60000000, lr=0.00025, gamma=0.99, C=20000, batch_size=32): env = wrappers.wrap(gym.make(GAMES[game])) num_actions = env.action_space.n Q1 = QNetwork(num_actions) Q2 = QNetwork(num_actions) Q2.load_state_dict(Q1.state_dict()) if torch.cuda.is_available(): Q1.cuda() Q2.cuda() epsilon = Epsilon(1, 0.05, 1000000) optimizer = torch.optim.Adam(Q1.parameters(), lr=lr) optimizer.zero_grad() state1 = env.reset() t, last_t, loss, episode, score = 0, 0, 0, 0, 0 last_ts, scores = datetime.now(), collections.deque(maxlen=100) while t < num_steps: qvalues = Q1(state1) if random() < epsilon(t): action = env.action_space.sample() else: action = qvalues.data.max(dim=1)[1][0] q = qvalues[0][action] state2, reward, done, _info = env.step(action) score += reward if not done: y = gamma * Q2(state2).detach().max(dim=1)[0][0] + reward state1 = state2 else: reward = FloatTensor([reward]) y = torch.autograd.Variable(reward, requires_grad=False) state1 = env.reset() scores.append(score) score = 0 episode += 1 loss += torch.nn.functional.smooth_l1_loss(q, y) t += 1 if done or t % batch_size == 0: loss.backward() optimizer.step() optimizer.zero_grad() loss = 0 if t % C == 0: Q2.load_state_dict(Q1.state_dict()) torch.save(Q1.state_dict(), 'qlearning_{}.pt'.format(game)) if t % 1000 == 0: ts = datetime.now() datestr = ts.strftime('%Y-%m-%dT%H:%M:%S.%f') avg = mean(scores) if scores else float('nan') steps_per_sec = (t - last_t) / (ts - last_ts).total_seconds() l = '{} step {} episode {} avg last 100 scores: {:.2f} ε: {:.2f}, steps/s: {:.0f}' print(l.format(datestr, t, episode, avg, epsilon(t), steps_per_sec)) last_t, last_ts = t, ts
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'), num_frames: ('int: Total number of 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, 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.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 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) self.memory = ReplayBuffer( self.buffer_size, (self.input_frames, self.input_dim, self.input_dim), self.batch_size) 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: # if normalization is applied to the image such as devision by 255, MUST be expressed 'state/255' below. state = torch.FloatTensor(state).to(self.device).unsqueeze(0) / 255 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, 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)) rewards = 0 dones = 0 for i in range(self.input_frames): 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[i] = self.processing_resize_and_gray(state) rewards += reward dones += int(done) return rewards, next_state, dones def get_init_state(self): state = self.env.reset() action = self.env.action_space.sample() _, state, _ = self.get_state(action, skipped_frame=0) return 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): batch = self.memory.batch_load() loss = self._compute_loss(batch) self.optimizer.zero_grad() loss.backward() self.optimizer.step() 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() for frame_idx in range(1, self.update_start + 1): _, action = self.select_action(state) reward, next_state, done = self.get_state(action, skipped_frame=0) self.store(state, action, reward, next_state, done) state = next_state if done: state = self.get_init_state() print("Done. Start learning..") history_store = [] for frame_idx in range(1, self.num_frames + 1): Qs, action = self.select_action(state) reward, next_state, done = self.get_state(action, skipped_frame=0) 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': wandb.log({ 'Score': score, 'loss(10 frames avg)': np.mean(losses[-10:]), 'Epsilon': self.epsilon }) print(score, end='\r') else: print(score, end='\r') score = 0 state = self.get_init_state() history_store = [] else: state = next_state self._epsilon_step() 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.35: 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) / 255 next_states = torch.FloatTensor(batch['next_states']).to( self.device) / 255 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) loss = F.smooth_l1_loss(current_q, target) return 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()
def train(env_name, seed=42, timesteps=1, epsilon_decay_last_step=1000, er_capacity=1e4, batch_size=16, lr=1e-3, gamma=1.0, update_target=16, exp_name='test', init_timesteps=100, save_every_steps=1e4, arch='nature', dueling=False, play_steps=2, n_jobs=2): """ Main training function. Calls the subprocesses to get experience and train the network. """ # Multiprocessing method mp.set_start_method('spawn') # Get PyTorch device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Set random seed for PyTorch torch.manual_seed(seed) torch.cuda.manual_seed(seed) # Create logger logger = Logger(exp_name, loggers=['tensorboard']) # Create the Q network _env = make_env(env_name, seed) net = QNetwork(_env.observation_space, _env.action_space, arch=arch, dueling=dueling).to(device) # Create the target network as a copy of the Q network target_net = copy.deepcopy(net) # Create buffer and optimizer buffer = ExperienceReplay(capacity=int(er_capacity)) optimizer = optim.Adam(net.parameters(), lr=lr) scheduler = StepLR(optimizer, step_size=LR_STEPS, gamma=0.99) # Multiprocessing queue obs_queue = mp.Queue(maxsize=n_jobs) transition_queue = mp.Queue(maxsize=n_jobs) workers, action_queues = [], [] for i in range(n_jobs): action_queue = mp.Queue(maxsize=1) _seed = seed + i * 1000 play_proc = mp.Process(target=play_func, args=(i, env_name, obs_queue, transition_queue, action_queue, _seed)) play_proc.start() workers.append(play_proc) action_queues.append(action_queue) # Vars to keep track of performances and time timestep = 0 current_reward, current_len = np.zeros(play_steps), np.zeros(play_steps, dtype=np.int64) current_time = [time.time() for _ in range(play_steps)] # Training loop while timestep < timesteps: # Compute the current epsilon epsilon = EPSILON_STOP + max(0, (EPSILON_START - EPSILON_STOP)*(epsilon_decay_last_step-timestep)/epsilon_decay_last_step) logger.log_kv('internals/epsilon', epsilon, timestep) # Gather observation N_STEPS ids, obs_batch = zip(*[obs_queue.get() for _ in range(play_steps)]) # Pre-process observation_batch for PyTorch obs_batch = torch.from_numpy(np.array(obs_batch)).to(device) # Select greedy action from policy, apply epsilon-greedy selection greedy_actions = net(obs_batch).argmax(dim=1).cpu().detach().numpy() probs = torch.rand(greedy_actions.shape) actions = np.where(probs < epsilon, _env.action_space.sample(), greedy_actions) # Send actions for id, action in zip(ids, actions): action_queues[id].put(action) # Add transitions to experience replay transitions = [transition_queue.get() for _ in range(play_steps)] buffer.pushTransitions(transitions) # Check if we need to update rewards, time and lengths _, _, _, reward, done, _ = zip(*transitions) current_reward += reward current_len += 1 for i, done in enumerate(done): if done: # Log quantities logger.log_kv('performance/return', current_reward[i], timestep) logger.log_kv('performance/length', current_len[i], timestep) logger.log_kv('performance/speed', current_len[i] / (time.time() - current_time[i]), timestep) # Reset counters current_reward[i] = 0.0 current_len[i] = 0 current_time[i] = time.time() # Update number of steps timestep += play_steps # Check if we are in the warm-up phase, otherwise go on with policy update if timestep < init_timesteps: continue # Learning rate upddate and log scheduler.step() logger.log_kv('internals/lr', scheduler.get_lr()[0], timestep) # Clear grads optimizer.zero_grad() # Get a batch from experience replay batch = buffer.sampleTransitions(batch_size) def batch_preprocess(batch_item): return torch.tensor(batch_item, dtype=(torch.long if isinstance(batch_item[0], np.int64) else None)).to(device) ids, states_batch, actions_batch, rewards_batch, done_batch, next_states_batch = map(batch_preprocess, zip(*batch)) # Compute the loss function state_action_values = net(states_batch).gather(1, actions_batch.unsqueeze(-1)).squeeze(-1) next_state_values = target_net(next_states_batch).max(1)[0] next_state_values[done_batch] = 0.0 expected_state_action_values = next_state_values.detach() * gamma + rewards_batch loss = F.mse_loss(state_action_values, expected_state_action_values) logger.log_kv('internals/loss', loss.item(), timestep) loss.backward() # Clip the gradients to avoid to abrupt changes (this is equivalent to Huber Loss) for param in net.parameters(): param.grad.data.clamp_(-1, 1) optimizer.step() if timestep % update_target == 0: target_net.load_state_dict(net.state_dict()) # Check if we need to save a checkpoint if timestep % save_every_steps == 0: torch.save(net.get_extended_state(), exp_name + '.pth') # Ending for i, worker in enumerate(workers): action_queues[i].put(None) worker.join()