def _initialize(self): """Initialize non-common things.""" self.per_beta = self.hyper_params.per_beta self.use_n_step = self.hyper_params.n_step > 1 if not self.args.test: # load demo replay memory with open(self.args.demo_path, "rb") as f: demos = pickle.load(f) if self.use_n_step: demos, demos_n_step = common_utils.get_n_step_info_from_demo( demos, self.hyper_params.n_step, self.hyper_params.gamma) # replay memory for multi-steps self.memory_n = ReplayBuffer( buffer_size=self.hyper_params.buffer_size, batch_size=self.hyper_params.batch_size, n_step=self.hyper_params.n_step, gamma=self.hyper_params.gamma, demo=demos_n_step, ) # replay memory self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, demo=demos, alpha=self.hyper_params.per_alpha, epsilon_d=self.hyper_params.per_eps_demo, )
def _initialize(self): """Initialize non-common things.""" if not self.args.test: # load demo replay memory demos = self._load_demos() if self.use_n_step: demos, demos_n_step = common_utils.get_n_step_info_from_demo( demos, self.hyper_params.n_step, self.hyper_params.gamma) self.memory_n = ReplayBuffer( buffer_size=self.hyper_params.buffer_size, n_step=self.hyper_params.n_step, gamma=self.hyper_params.gamma, demo=demos_n_step, ) # replay memory self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, demo=demos, alpha=self.hyper_params.per_alpha, epsilon_d=self.hyper_params.per_eps_demo, )
def _initialize(self): """Initialize non-common things.""" self.per_beta = self.hyper_params.per_beta if not self.args.test: # replay memory self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, alpha=self.hyper_params.per_alpha, )
def _initialize(self): """Initialize non-common things.""" if not self.args.test: # replay memory for a single step self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, alpha=self.hyper_params.per_alpha, ) # replay memory for multi-steps if self.use_n_step: self.memory_n = ReplayBuffer( self.hyper_params.buffer_size, n_step=self.hyper_params.n_step, gamma=self.hyper_params.gamma, )
class DQfDAgent(DQNAgent): """DQN interacting with environment. Attribute: memory (PrioritizedReplayBuffer): replay memory """ # pylint: disable=attribute-defined-outside-init def _initialize(self): """Initialize non-common things.""" if not self.args.test: # load demo replay memory demos = self._load_demos() if self.use_n_step: demos, demos_n_step = common_utils.get_n_step_info_from_demo( demos, self.hyper_params.n_step, self.hyper_params.gamma) self.memory_n = ReplayBuffer( buffer_size=self.hyper_params.buffer_size, n_step=self.hyper_params.n_step, gamma=self.hyper_params.gamma, demo=demos_n_step, ) # replay memory self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, demo=demos, alpha=self.hyper_params.per_alpha, epsilon_d=self.hyper_params.per_eps_demo, ) def _load_demos(self) -> list: """Load expert's demonstrations.""" # load demo replay memory with open(self.args.demo_path, "rb") as f: demos = pickle.load(f) return demos def update_model(self) -> Tuple[torch.Tensor, ...]: """Train the model after each episode.""" experiences_1 = self.memory.sample() weights, indices, eps_d = experiences_1[-3:] actions = experiences_1[1] # 1 step loss gamma = self.hyper_params.gamma dq_loss_element_wise, q_values = self._get_dqn_loss( experiences_1, gamma) dq_loss = torch.mean(dq_loss_element_wise * weights) # n step loss if self.use_n_step: experiences_n = self.memory_n.sample(indices) gamma = self.hyper_params.gamma**self.hyper_params.n_step dq_loss_n_element_wise, q_values_n = self._get_dqn_loss( experiences_n, gamma) # to update loss and priorities q_values = 0.5 * (q_values + q_values_n) dq_loss_element_wise += dq_loss_n_element_wise * self.hyper_params.lambda1 dq_loss = torch.mean(dq_loss_element_wise * weights) # supervised loss using demo for only demo transitions demo_idxs = np.where(eps_d != 0.0) n_demo = demo_idxs[0].size if n_demo != 0: # if 1 or more demos are sampled # get margin for each demo transition action_idxs = actions[demo_idxs].long() margin = torch.ones(q_values.size()) * self.hyper_params.margin margin[demo_idxs, action_idxs] = 0.0 # demo actions have 0 margins margin = margin.to(device) # calculate supervised loss demo_q_values = q_values[demo_idxs, action_idxs].squeeze() supervised_loss = torch.max(q_values + margin, dim=-1)[0] supervised_loss = supervised_loss[demo_idxs] - demo_q_values supervised_loss = torch.mean( supervised_loss) * self.hyper_params.lambda2 else: # no demo sampled supervised_loss = torch.zeros(1, device=device) # q_value regularization q_regular = torch.norm(q_values, 2).mean() * self.hyper_params.w_q_reg # total loss loss = dq_loss + supervised_loss + q_regular # train dqn self.dqn_optim.zero_grad() loss.backward() clip_grad_norm_(self.dqn.parameters(), self.hyper_params.gradient_clip) self.dqn_optim.step() # update target networks common_utils.soft_update(self.dqn, self.dqn_target, self.hyper_params.tau) # update priorities in PER loss_for_prior = dq_loss_element_wise.detach().cpu().numpy().squeeze() new_priorities = loss_for_prior + self.hyper_params.per_eps new_priorities += eps_d self.memory.update_priorities(indices, new_priorities) # increase beta fraction = min(float(self.i_episode) / self.args.episode_num, 1.0) self.per_beta: float = self.per_beta + fraction * (1.0 - self.per_beta) if self.hyper_params.use_noisy_net: self.dqn.reset_noise() self.dqn_target.reset_noise() return ( loss.item(), dq_loss.item(), supervised_loss.item(), q_values.mean().item(), n_demo, ) def write_log(self, log_value: tuple): """Write log about loss and score""" i, avg_loss, score, avg_time_cost = log_value print( "[INFO] episode %d, episode step: %d, total step: %d, total score: %f\n" "epsilon: %f, total loss: %f, dq loss: %f, supervised loss: %f\n" "avg q values: %f, demo num in minibatch: %d (spent %.6f sec/step)\n" % ( i, self.episode_step, self.total_step, score, self.epsilon, avg_loss[0], avg_loss[1], avg_loss[2], avg_loss[3], avg_loss[4], avg_time_cost, )) if self.args.log: wandb.log({ "score": score, "epsilon": self.epsilon, "total loss": avg_loss[0], "dq loss": avg_loss[1], "supervised loss": avg_loss[2], "avg q values": avg_loss[3], "demo num in minibatch": avg_loss[4], "time per each step": avg_time_cost, }) def pretrain(self): """Pretraining steps.""" pretrain_loss = list() pretrain_step = self.hyper_params.pretrain_step print("[INFO] Pre-Train %d step." % pretrain_step) for i_step in range(1, pretrain_step + 1): t_begin = time.time() loss = self.update_model() t_end = time.time() pretrain_loss.append(loss) # for logging # logging if i_step == 1 or i_step % 100 == 0: avg_loss = np.vstack(pretrain_loss).mean(axis=0) pretrain_loss.clear() log_value = (0, avg_loss, 0.0, t_end - t_begin) self.write_log(log_value) print("[INFO] Pre-Train Complete!\n")
class DQNAgent(Agent): """DQN interacting with environment. Attribute: env (gym.Env): openAI Gym environment args (argparse.Namespace): arguments including hyperparameters and training settings hyper_params (ConfigDict): hyper-parameters network_cfg (ConfigDict): config of network for training agent optim_cfg (ConfigDict): config of optimizer state_dim (int): state size of env action_dim (int): action size of env memory (PrioritizedReplayBuffer): replay memory dqn (nn.Module): actor model to select actions dqn_target (nn.Module): target actor model to select actions dqn_optim (Optimizer): optimizer for training actor curr_state (np.ndarray): temporary storage of the current state total_step (int): total step number episode_step (int): step number of the current episode i_episode (int): current episode number epsilon (float): parameter for epsilon greedy policy n_step_buffer (deque): n-size buffer to calculate n-step returns per_beta (float): beta parameter for prioritized replay buffer use_conv (bool): whether or not to use convolution layer use_n_step (bool): whether or not to use n-step returns """ def __init__( self, env: gym.Env, args: argparse.Namespace, log_cfg: ConfigDict, hyper_params: ConfigDict, network_cfg: ConfigDict, optim_cfg: ConfigDict, ): """Initialize.""" Agent.__init__(self, env, args, log_cfg) self.curr_state = np.zeros(1) self.episode_step = 0 self.total_step = 0 self.i_episode = 0 self.hyper_params = hyper_params self.network_cfg = network_cfg self.optim_cfg = optim_cfg self.state_dim = self.env.observation_space.shape self.action_dim = self.env.action_space.n self.per_beta = hyper_params.per_beta self.use_conv = len(self.state_dim) > 1 self.use_n_step = hyper_params.n_step > 1 if hyper_params.use_noisy_net: self.max_epsilon = 0.0 self.min_epsilon = 0.0 self.epsilon = 0.0 else: self.max_epsilon = hyper_params.max_epsilon self.min_epsilon = hyper_params.min_epsilon self.epsilon = hyper_params.max_epsilon self._initialize() self._init_network() # pylint: disable=attribute-defined-outside-init def _initialize(self): """Initialize non-common things.""" if not self.args.test: # replay memory for a single step self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, alpha=self.hyper_params.per_alpha, ) # replay memory for multi-steps if self.use_n_step: self.memory_n = ReplayBuffer( self.hyper_params.buffer_size, n_step=self.hyper_params.n_step, gamma=self.hyper_params.gamma, ) # pylint: disable=attribute-defined-outside-init def _init_network(self): """Initialize networks and optimizers.""" if self.use_conv: # create CNN self.dqn = dqn_utils.get_cnn_model(self.hyper_params, self.action_dim, self.state_dim, self.network_cfg) self.dqn_target = dqn_utils.get_cnn_model(self.hyper_params, self.action_dim, self.state_dim, self.network_cfg) else: # create FC fc_input_size = self.state_dim[0] self.dqn = dqn_utils.get_fc_model( self.hyper_params, fc_input_size, self.action_dim, self.network_cfg.hidden_sizes, ) self.dqn_target = dqn_utils.get_fc_model( self.hyper_params, fc_input_size, self.action_dim, self.network_cfg.hidden_sizes, ) self.dqn_target.load_state_dict(self.dqn.state_dict()) # create optimizer self.dqn_optim = optim.Adam( self.dqn.parameters(), lr=self.optim_cfg.lr_dqn, weight_decay=self.optim_cfg.weight_decay, eps=self.optim_cfg.adam_eps, ) # load the optimizer and model parameters if self.args.load_from is not None: self.load_params(self.args.load_from) def select_action(self, state: np.ndarray) -> np.ndarray: """Select an action from the input space.""" self.curr_state = state # epsilon greedy policy # pylint: disable=comparison-with-callable if not self.args.test and self.epsilon > np.random.random(): selected_action = np.array(self.env.action_space.sample()) else: state = self._preprocess_state(state) selected_action = self.dqn(state).argmax() selected_action = selected_action.detach().cpu().numpy() return selected_action # pylint: disable=no-self-use def _preprocess_state(self, state: np.ndarray) -> torch.Tensor: """Preprocess state so that actor selects an action.""" state = torch.FloatTensor(state).to(device) return state def step(self, action: np.ndarray) -> Tuple[np.ndarray, np.float64, bool, dict]: """Take an action and return the response of the env.""" next_state, reward, done, info = self.env.step(action) if not self.args.test: # if the last state is not a terminal state, store done as false done_bool = (False if self.episode_step == self.args.max_episode_steps else done) transition = (self.curr_state, action, reward, next_state, done_bool) self._add_transition_to_memory(transition) return next_state, reward, done, info def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): """Add 1 step and n step transitions to memory.""" # add n-step transition if self.use_n_step: transition = self.memory_n.add(transition) # add a single step transition # if transition is not an empty tuple if transition: self.memory.add(transition) def _get_dqn_loss(self, experiences: Tuple[torch.Tensor, ...], gamma: float) -> Tuple[torch.Tensor, torch.Tensor]: """Return element-wise dqn loss and Q-values.""" if self.hyper_params.use_dist_q == "IQN": return dqn_utils.calculate_iqn_loss( model=self.dqn, target_model=self.dqn_target, experiences=experiences, gamma=gamma, batch_size=self.hyper_params.batch_size, n_tau_samples=self.hyper_params.n_tau_samples, n_tau_prime_samples=self.hyper_params.n_tau_prime_samples, kappa=self.hyper_params.kappa, ) elif self.hyper_params.use_dist_q == "C51": return dqn_utils.calculate_c51_loss( model=self.dqn, target_model=self.dqn_target, experiences=experiences, gamma=gamma, batch_size=self.hyper_params.batch_size, v_min=self.hyper_params.v_min, v_max=self.hyper_params.v_max, atom_size=self.hyper_params.atoms, ) else: return dqn_utils.calculate_dqn_loss( model=self.dqn, target_model=self.dqn_target, experiences=experiences, gamma=gamma, ) def update_model(self) -> Tuple[torch.Tensor, ...]: """Train the model after each episode.""" # 1 step loss experiences_1 = self.memory.sample(self.per_beta) weights, indices = experiences_1[-3:-1] gamma = self.hyper_params.gamma dq_loss_element_wise, q_values = self._get_dqn_loss( experiences_1, gamma) dq_loss = torch.mean(dq_loss_element_wise * weights) # n step loss if self.use_n_step: experiences_n = self.memory_n.sample(indices) gamma = self.hyper_params.gamma**self.hyper_params.n_step dq_loss_n_element_wise, q_values_n = self._get_dqn_loss( experiences_n, gamma) # to update loss and priorities q_values = 0.5 * (q_values + q_values_n) dq_loss_element_wise += dq_loss_n_element_wise * self.hyper_params.w_n_step dq_loss = torch.mean(dq_loss_element_wise * weights) # q_value regularization q_regular = torch.norm(q_values, 2).mean() * self.hyper_params.w_q_reg # total loss loss = dq_loss + q_regular self.dqn_optim.zero_grad() loss.backward() clip_grad_norm_(self.dqn.parameters(), self.hyper_params.gradient_clip) self.dqn_optim.step() # update target networks common_utils.soft_update(self.dqn, self.dqn_target, self.hyper_params.tau) # update priorities in PER loss_for_prior = dq_loss_element_wise.detach().cpu().numpy() new_priorities = loss_for_prior + self.hyper_params.per_eps self.memory.update_priorities(indices, new_priorities) # increase beta fraction = min(float(self.i_episode) / self.args.episode_num, 1.0) self.per_beta = self.per_beta + fraction * (1.0 - self.per_beta) if self.hyper_params.use_noisy_net: self.dqn.reset_noise() self.dqn_target.reset_noise() return loss.item(), q_values.mean().item() def load_params(self, path: str): """Load model and optimizer parameters.""" Agent.load_params(self, path) params = torch.load(path) self.dqn.load_state_dict(params["dqn_state_dict"]) self.dqn_target.load_state_dict(params["dqn_target_state_dict"]) self.dqn_optim.load_state_dict(params["dqn_optim_state_dict"]) print("[INFO] loaded the model and optimizer from", path) def save_params(self, n_episode: int): # type: ignore """Save model and optimizer parameters.""" params = { "dqn_state_dict": self.dqn.state_dict(), "dqn_target_state_dict": self.dqn_target.state_dict(), "dqn_optim_state_dict": self.dqn_optim.state_dict(), } Agent.save_params(self, params, n_episode) def write_log(self, log_value: tuple): """Write log about loss and score""" i, loss, score, avg_time_cost = log_value print( "[INFO] episode %d, episode step: %d, total step: %d, total score: %f\n" "epsilon: %f, loss: %f, avg q-value: %f (spent %.6f sec/step)\n" % ( i, self.episode_step, self.total_step, score, self.epsilon, loss[0], loss[1], avg_time_cost, )) if self.args.log: wandb.log({ "score": score, "epsilon": self.epsilon, "dqn loss": loss[0], "avg q values": loss[1], "time per each step": avg_time_cost, }) # pylint: disable=no-self-use, unnecessary-pass def pretrain(self): """Pretraining steps.""" pass def train(self): """Train the agent.""" # logger if self.args.log: self.set_wandb() # wandb.watch([self.dqn], log="parameters") # pre-training if needed self.pretrain() for self.i_episode in range(1, self.args.episode_num + 1): state = self.env.reset() self.episode_step = 0 losses = list() done = False score = 0 t_begin = time.time() while not done: if self.args.render and self.i_episode >= self.args.render_after: self.env.render() action = self.select_action(state) next_state, reward, done, _ = self.step(action) self.total_step += 1 self.episode_step += 1 if len(self.memory) >= self.hyper_params.update_starts_from: if self.total_step % self.hyper_params.train_freq == 0: for _ in range(self.hyper_params.multiple_update): loss = self.update_model() losses.append(loss) # for logging # decrease epsilon self.epsilon = max( self.epsilon - (self.max_epsilon - self.min_epsilon) * self.hyper_params.epsilon_decay, self.min_epsilon, ) state = next_state score += reward t_end = time.time() avg_time_cost = (t_end - t_begin) / self.episode_step if losses: avg_loss = np.vstack(losses).mean(axis=0) log_value = (self.i_episode, avg_loss, score, avg_time_cost) self.write_log(log_value) if self.i_episode % self.args.save_period == 0: self.save_params(self.i_episode) self.interim_test() # termination self.env.close() self.save_params(self.i_episode) self.interim_test()
class SACfDAgent(SACAgent): """SAC agent interacting with environment. Attrtibutes: memory (PrioritizedReplayBuffer): replay memory beta (float): beta parameter for prioritized replay buffer use_n_step (bool): whether or not to use n-step returns """ # pylint: disable=attribute-defined-outside-init def _initialize(self): """Initialize non-common things.""" self.per_beta = self.hyper_params.per_beta self.use_n_step = self.hyper_params.n_step > 1 if not self.args.test: # load demo replay memory with open(self.args.demo_path, "rb") as f: demos = pickle.load(f) if self.use_n_step: demos, demos_n_step = common_utils.get_n_step_info_from_demo( demos, self.hyper_params.n_step, self.hyper_params.gamma) # replay memory for multi-steps self.memory_n = ReplayBuffer( buffer_size=self.hyper_params.buffer_size, batch_size=self.hyper_params.batch_size, n_step=self.hyper_params.n_step, gamma=self.hyper_params.gamma, demo=demos_n_step, ) # replay memory self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, demo=demos, alpha=self.hyper_params.per_alpha, epsilon_d=self.hyper_params.per_eps_demo, ) def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): """Add 1 step and n step transitions to memory.""" # add n-step transition if self.use_n_step: transition = self.memory_n.add(transition) # add a single step transition # if transition is not an empty tuple if transition: self.memory.add(transition) # pylint: disable=too-many-statements def update_model(self) -> Tuple[torch.Tensor, ...]: """Train the model after each episode.""" self.update_step += 1 experiences = self.memory.sample(self.per_beta) ( states, actions, rewards, next_states, dones, weights, indices, eps_d, ) = experiences new_actions, log_prob, pre_tanh_value, mu, std = self.actor(states) # train alpha if self.hyper_params.auto_entropy_tuning: alpha_loss = torch.mean( (-self.log_alpha * (log_prob + self.target_entropy).detach()) * weights) self.alpha_optim.zero_grad() alpha_loss.backward() self.alpha_optim.step() alpha = self.log_alpha.exp() else: alpha_loss = torch.zeros(1) alpha = self.hyper_params.w_entropy # Q function loss masks = 1 - dones gamma = self.hyper_params.gamma states_actions = torch.cat((states, actions), dim=-1) q_1_pred = self.qf_1(states_actions) q_2_pred = self.qf_2(states_actions) v_target = self.vf_target(next_states) q_target = rewards + self.hyper_params.gamma * v_target * masks qf_1_loss = torch.mean((q_1_pred - q_target.detach()).pow(2) * weights) qf_2_loss = torch.mean((q_2_pred - q_target.detach()).pow(2) * weights) if self.use_n_step: experiences_n = self.memory_n.sample(indices) _, _, rewards, next_states, dones = experiences_n gamma = gamma**self.hyper_params.n_step masks = 1 - dones v_target = self.vf_target(next_states) q_target = rewards + gamma * v_target * masks qf_1_loss_n = torch.mean( (q_1_pred - q_target.detach()).pow(2) * weights) qf_2_loss_n = torch.mean( (q_2_pred - q_target.detach()).pow(2) * weights) # to update loss and priorities qf_1_loss = qf_1_loss + qf_1_loss_n * self.hyper_params.lambda1 qf_2_loss = qf_2_loss + qf_2_loss_n * self.hyper_params.lambda1 # V function loss states_actions = torch.cat((states, new_actions), dim=-1) v_pred = self.vf(states) q_pred = torch.min(self.qf_1(states_actions), self.qf_2(states_actions)) v_target = (q_pred - alpha * log_prob).detach() vf_loss_element_wise = (v_pred - v_target).pow(2) vf_loss = torch.mean(vf_loss_element_wise * weights) # train Q functions self.qf_1_optim.zero_grad() qf_1_loss.backward() self.qf_1_optim.step() self.qf_2_optim.zero_grad() qf_2_loss.backward() self.qf_2_optim.step() # train V function self.vf_optim.zero_grad() vf_loss.backward() self.vf_optim.step() if self.update_step % self.hyper_params.policy_update_freq == 0: # actor loss advantage = q_pred - v_pred.detach() actor_loss_element_wise = alpha * log_prob - advantage actor_loss = torch.mean(actor_loss_element_wise * weights) # regularization mean_reg = self.hyper_params.w_mean_reg * mu.pow(2).mean() std_reg = self.hyper_params.w_std_reg * std.pow(2).mean() pre_activation_reg = self.hyper_params.w_pre_activation_reg * ( pre_tanh_value.pow(2).sum(dim=-1).mean()) actor_reg = mean_reg + std_reg + pre_activation_reg # actor loss + regularization actor_loss += actor_reg # train actor self.actor_optim.zero_grad() actor_loss.backward() self.actor_optim.step() # update target networks common_utils.soft_update(self.vf, self.vf_target, self.hyper_params.tau) # update priorities new_priorities = vf_loss_element_wise new_priorities += self.hyper_params.lambda3 * actor_loss_element_wise.pow( 2) new_priorities += self.hyper_params.per_eps new_priorities = new_priorities.data.cpu().numpy().squeeze() new_priorities += eps_d self.memory.update_priorities(indices, new_priorities) # increase beta fraction = min(float(self.i_episode) / self.args.episode_num, 1.0) self.per_beta = self.per_beta + fraction * (1.0 - self.per_beta) else: actor_loss = torch.zeros(1) return ( actor_loss.item(), qf_1_loss.item(), qf_2_loss.item(), vf_loss.item(), alpha_loss.item(), ) def pretrain(self): """Pretraining steps.""" pretrain_loss = list() pretrain_step = self.hyper_params.pretrain_step print("[INFO] Pre-Train %d steps." % pretrain_step) for i_step in range(1, pretrain_step + 1): t_begin = time.time() loss = self.update_model() t_end = time.time() pretrain_loss.append(loss) # for logging # logging if i_step == 1 or i_step % 100 == 0: avg_loss = np.vstack(pretrain_loss).mean(axis=0) pretrain_loss.clear() log_value = ( 0, avg_loss, 0, self.hyper_params.policy_update_freq, t_end - t_begin, ) self.write_log(log_value) print("[INFO] Pre-Train Complete!\n")
class DDPGfDAgent(DDPGAgent): """ActorCritic interacting with environment. Attributes: memory (PrioritizedReplayBuffer): replay memory per_beta (float): beta parameter for prioritized replay buffer use_n_step (bool): whether or not to use n-step returns """ # pylint: disable=attribute-defined-outside-init def _initialize(self): """Initialize non-common things.""" self.per_beta = self.hyper_params.per_beta self.use_n_step = self.hyper_params.n_step > 1 if not self.args.test: # load demo replay memory with open(self.args.demo_path, "rb") as f: demos = pickle.load(f) if self.use_n_step: demos, demos_n_step = common_utils.get_n_step_info_from_demo( demos, self.hyper_params.n_step, self.hyper_params.gamma ) # replay memory for multi-steps self.memory_n = ReplayBuffer( buffer_size=self.hyper_params.buffer_size, batch_size=self.hyper_params.batch_size, n_step=self.hyper_params.n_step, gamma=self.hyper_params.gamma, demo=demos_n_step, ) # replay memory for a single step self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, demo=demos, alpha=self.hyper_params.per_alpha, epsilon_d=self.hyper_params.per_eps_demo, ) def _add_transition_to_memory(self, transition: Tuple[np.ndarray, ...]): """Add 1 step and n step transitions to memory.""" # add n-step transition if self.use_n_step: transition = self.memory_n.add(transition) # add a single step transition # if transition is not an empty tuple if transition: self.memory.add(transition) def _get_critic_loss( self, experiences: Tuple[torch.Tensor, ...], gamma: float ) -> torch.Tensor: """Return element-wise critic loss.""" states, actions, rewards, next_states, dones = experiences[:5] # G_t = r + gamma * v(s_{t+1}) if state != Terminal # = r otherwise masks = 1 - dones next_actions = self.actor_target(next_states) next_states_actions = torch.cat((next_states, next_actions), dim=-1) next_values = self.critic_target(next_states_actions) curr_returns = rewards + gamma * next_values * masks curr_returns = curr_returns.to(device).detach() # train critic values = self.critic(torch.cat((states, actions), dim=-1)) critic_loss_element_wise = (values - curr_returns).pow(2) return critic_loss_element_wise def update_model(self) -> Tuple[torch.Tensor, ...]: """Train the model after each episode.""" experiences_1 = self.memory.sample(self.per_beta) states, actions = experiences_1[:2] weights, indices, eps_d = experiences_1[-3:] gamma = self.hyper_params.gamma # train critic gradient_clip_ac = self.hyper_params.gradient_clip_ac gradient_clip_cr = self.hyper_params.gradient_clip_cr critic_loss_element_wise = self._get_critic_loss(experiences_1, gamma) critic_loss = torch.mean(critic_loss_element_wise * weights) if self.use_n_step: experiences_n = self.memory_n.sample(indices) gamma = gamma ** self.hyper_params.n_step critic_loss_n_element_wise = self._get_critic_loss(experiences_n, gamma) # to update loss and priorities critic_loss_element_wise += ( critic_loss_n_element_wise * self.hyper_params.lambda1 ) critic_loss = torch.mean(critic_loss_element_wise * weights) self.critic_optim.zero_grad() critic_loss.backward() nn.utils.clip_grad_norm_(self.critic.parameters(), gradient_clip_cr) self.critic_optim.step() # train actor actions = self.actor(states) actor_loss_element_wise = -self.critic(torch.cat((states, actions), dim=-1)) actor_loss = torch.mean(actor_loss_element_wise * weights) self.actor_optim.zero_grad() actor_loss.backward() nn.utils.clip_grad_norm_(self.actor.parameters(), gradient_clip_ac) self.actor_optim.step() # update target networks common_utils.soft_update(self.actor, self.actor_target, self.hyper_params.tau) common_utils.soft_update(self.critic, self.critic_target, self.hyper_params.tau) # update priorities new_priorities = critic_loss_element_wise new_priorities += self.hyper_params.lambda3 * actor_loss_element_wise.pow(2) new_priorities += self.hyper_params.per_eps new_priorities = new_priorities.data.cpu().numpy().squeeze() new_priorities += eps_d self.memory.update_priorities(indices, new_priorities) # increase beta fraction = min(float(self.i_episode) / self.args.episode_num, 1.0) self.per_beta = self.per_beta + fraction * (1.0 - self.per_beta) return actor_loss.item(), critic_loss.item() def pretrain(self): """Pretraining steps.""" pretrain_loss = list() pretrain_step = self.hyper_params.pretrain_step print("[INFO] Pre-Train %d step." % pretrain_step) for i_step in range(1, pretrain_step + 1): t_begin = time.time() loss = self.update_model() t_end = time.time() pretrain_loss.append(loss) # for logging # logging if i_step == 1 or i_step % 100 == 0: avg_loss = np.vstack(pretrain_loss).mean(axis=0) pretrain_loss.clear() log_value = (0, avg_loss, 0, t_end - t_begin) self.write_log(log_value) print("[INFO] Pre-Train Complete!\n")
class PERDDPGAgent(DDPGAgent): """ActorCritic interacting with environment. Attributes: memory (PrioritizedReplayBuffer): replay memory per_beta (float): beta parameter for prioritized replay buffer """ # pylint: disable=attribute-defined-outside-init def _initialize(self): """Initialize non-common things.""" self.per_beta = self.hyper_params.per_beta if not self.args.test: # replay memory self.memory = PrioritizedReplayBuffer( self.hyper_params.buffer_size, self.hyper_params.batch_size, alpha=self.hyper_params.per_alpha, ) def update_model(self) -> Tuple[torch.Tensor, ...]: """Train the model after each episode.""" experiences = self.memory.sample(self.per_beta) states, actions, rewards, next_states, dones, weights, indices, _ = experiences # G_t = r + gamma * v(s_{t+1}) if state != Terminal # = r otherwise masks = 1 - dones next_actions = self.actor_target(next_states) next_values = self.critic_target(torch.cat((next_states, next_actions), dim=-1)) curr_returns = rewards + self.hyper_params.gamma * next_values * masks curr_returns = curr_returns.to(device).detach() # train critic gradient_clip_ac = self.hyper_params.gradient_clip_ac gradient_clip_cr = self.hyper_params.gradient_clip_cr values = self.critic(torch.cat((states, actions), dim=-1)) critic_loss_element_wise = (values - curr_returns).pow(2) critic_loss = torch.mean(critic_loss_element_wise * weights) self.critic_optim.zero_grad() critic_loss.backward() nn.utils.clip_grad_norm_(self.critic.parameters(), gradient_clip_cr) self.critic_optim.step() # train actor actions = self.actor(states) actor_loss_element_wise = -self.critic(torch.cat((states, actions), dim=-1)) actor_loss = torch.mean(actor_loss_element_wise * weights) self.actor_optim.zero_grad() actor_loss.backward() nn.utils.clip_grad_norm_(self.actor.parameters(), gradient_clip_ac) self.actor_optim.step() # update target networks common_utils.soft_update(self.actor, self.actor_target, self.hyper_params.tau) common_utils.soft_update(self.critic, self.critic_target, self.hyper_params.tau) # update priorities in PER new_priorities = critic_loss_element_wise new_priorities = new_priorities.data.cpu().numpy() + self.hyper_params.per_eps self.memory.update_priorities(indices, new_priorities) # increase beta fraction = min(float(self.i_episode) / self.args.episode_num, 1.0) self.per_beta = self.per_beta + fraction * (1.0 - self.per_beta) return actor_loss.item(), critic_loss.item()