(self, replay_buffer, batch_size, current_map_level)
| 255 | return action, log_prob, z, batch_mu, batch_log_sigma |
| 256 | |
| 257 | def update(self, replay_buffer, batch_size, current_map_level): |
| 258 | for _ in range(self.args.gradient_steps): |
| 259 | |
| 260 | self.tem = self.initial_tem / (1.0 + self.tem_decay_rate * self.num_training_map_level) |
| 261 | |
| 262 | # Sample a batch from replaybuffer |
| 263 | (s, a, r, s_, d) = replay_buffer.sample_batch(batch_size, self.args.keys_num) |
| 264 | bn_s = torch.Tensor(s).float().to(device) |
| 265 | bn_a = torch.Tensor(a).to(device) |
| 266 | bn_r = torch.Tensor(r).to(device) |
| 267 | bn_s_ = torch.Tensor(s_).float().to(device) |
| 268 | bn_d = torch.Tensor(d).float().to(device) |
| 269 | |
| 270 | bn_r = bn_r.reshape(-1, self.args.keys_num, 1)[:, -1, :] |
| 271 | bn_d = bn_d.reshape(-1, self.args.keys_num, 1)[:, -1, :] |
| 272 | target_value = self.Target_value_net(bn_s_) |
| 273 | next_q_value = bn_r + (1 - bn_d) * self.args.gamma * target_value |
| 274 | |
| 275 | excepted_value = self.value_net(bn_s) |
| 276 | excepted_Q1 = self.Q_net1(bn_s, bn_a) |
| 277 | excepted_Q2 = self.Q_net2(bn_s, bn_a) |
| 278 | |
| 279 | sample_action, log_prob, *_ = self.evaluate(bn_s) |
| 280 | log_prob = log_prob[:, -1, :] |
| 281 | excepted_new_Q = torch.min(self.Q_net1(bn_s, sample_action), self.Q_net2(bn_s, sample_action)) |
| 282 | next_value = excepted_new_Q - self.tem * log_prob #J_V |
| 283 | |
| 284 | # !!!Note that the actions are sampled according to the current policy, |
| 285 | # instead of replay buffer. (From original paper) |
| 286 | V_loss = self.value_criterion(excepted_value, next_value.detach()).mean() # J_V |
| 287 | |
| 288 | # Dual Q net |
| 289 | Q1_loss = self.Q1_criterion(excepted_Q1, next_q_value.detach()).mean() # J_Q |
| 290 | Q2_loss = self.Q2_criterion(excepted_Q2, next_q_value.detach()).mean() |
| 291 | |
| 292 | self.writer.add_scalar('Loss/V_loss', V_loss, global_step=self.num_training) |
| 293 | self.writer.add_scalar('Loss/Q1_loss', Q1_loss, global_step=self.num_training) |
| 294 | self.writer.add_scalar('Loss/Q2_loss', Q2_loss, global_step=self.num_training) |
| 295 | |
| 296 | # mini batch gradient descent |
| 297 | self.value_optimizer.zero_grad() |
| 298 | V_loss.backward(retain_graph=True) |
| 299 | nn.utils.clip_grad_norm_(self.value_net.parameters(), 0.5) |
| 300 | self.value_optimizer.step() |
| 301 | |
| 302 | self.Q1_optimizer.zero_grad() |
| 303 | Q1_loss.backward(retain_graph = True) |
| 304 | nn.utils.clip_grad_norm_(self.Q_net1.parameters(), 0.5) |
| 305 | self.Q1_optimizer.step() |
| 306 | |
| 307 | self.Q2_optimizer.zero_grad() |
| 308 | Q2_loss.backward(retain_graph = True) |
| 309 | nn.utils.clip_grad_norm_(self.Q_net2.parameters(), 0.5) |
| 310 | self.Q2_optimizer.step() |
| 311 | if current_map_level != self.last_map_level: |
| 312 | self.num_training_map_level = 0 |
| 313 | |
| 314 | if self.num_training % self.args.policy_update_interval == 0: |
no test coverage detected