| 9 | |
| 10 | |
| 11 | def _log(config, step_count, log_data, model, replay_buffer, lr, shared_storage, summary_writer, vis_result): |
| 12 | loss_data, td_data, priority_data = log_data |
| 13 | total_loss, weighted_loss, loss, reg_loss, policy_loss, value_prefix_loss, value_loss, consistency_loss = loss_data |
| 14 | if vis_result: |
| 15 | new_priority, target_value_prefix, target_value, trans_target_value_prefix, trans_target_value, target_value_prefix_phi, target_value_phi, \ |
| 16 | pred_value_prefix, pred_value, target_policies, predicted_policies, state_lst, other_loss, other_log, other_dist = td_data |
| 17 | batch_weights, batch_indices = priority_data |
| 18 | |
| 19 | replay_episodes_collected, replay_buffer_size, priorities, total_num, worker_logs = ray.get([ |
| 20 | replay_buffer.episodes_collected.remote(), replay_buffer.size.remote(), |
| 21 | replay_buffer.get_priorities.remote(), replay_buffer.get_total_len.remote(), |
| 22 | shared_storage.get_worker_logs.remote()]) |
| 23 | |
| 24 | worker_ori_reward, worker_reward, worker_reward_max, worker_eps_len, worker_eps_len_max, test_counter, test_dict, temperature, visit_entropy, priority_self_play, distributions = worker_logs |
| 25 | |
| 26 | _msg = '#{:<10} Total Loss: {:<8.3f} [weighted Loss:{:<8.3f} Policy Loss: {:<8.3f} Value Loss: {:<8.3f} ' \ |
| 27 | 'Reward Sum Loss: {:<8.3f} Consistency Loss: {:<8.3f} ] ' \ |
| 28 | 'Replay Episodes Collected: {:<10d} Buffer Size: {:<10d} Transition Number: {:<8.3f}k ' \ |
| 29 | 'Batch Size: {:<10d} Lr: {:<8.3f}' |
| 30 | _msg = _msg.format(step_count, total_loss, weighted_loss, policy_loss, value_loss, value_prefix_loss, consistency_loss, |
| 31 | replay_episodes_collected, replay_buffer_size, total_num / 1000, config.batch_size, lr) |
| 32 | train_logger.info(_msg) |
| 33 | |
| 34 | if test_dict is not None: |
| 35 | mean_score = np.mean(test_dict['mean_score']) |
| 36 | max_score = np.mean(test_dict['max_score']) |
| 37 | min_score = np.mean(test_dict['min_score']) |
| 38 | std_score = np.mean(test_dict['std_score']) |
| 39 | test_msg = '#{:<10} Test Mean Score of {}: {:<10} (max: {:<10}, min:{:<10}, std: {:<10})' \ |
| 40 | ''.format(test_counter, config.env_name, mean_score, max_score, min_score, std_score) |
| 41 | test_logger.info(test_msg) |
| 42 | |
| 43 | if summary_writer is not None: |
| 44 | if config.debug: |
| 45 | for name, W in model.named_parameters(): |
| 46 | summary_writer.add_histogram('after_grad_clip' + '/' + name + '_grad', W.grad.data.cpu().numpy(), |
| 47 | step_count) |
| 48 | summary_writer.add_histogram('network_weights' + '/' + name, W.data.cpu().numpy(), step_count) |
| 49 | pass |
| 50 | tag = 'Train' |
| 51 | if vis_result: |
| 52 | summary_writer.add_histogram('{}_replay_data/replay_buffer_priorities'.format(tag), |
| 53 | priorities, |
| 54 | step_count) |
| 55 | summary_writer.add_histogram('{}_replay_data/batch_weight'.format(tag), batch_weights, step_count) |
| 56 | summary_writer.add_histogram('{}_replay_data/batch_indices'.format(tag), batch_indices, step_count) |
| 57 | target_value_prefix = target_value_prefix.flatten() |
| 58 | pred_value_prefix = pred_value_prefix.flatten() |
| 59 | target_value = target_value.flatten() |
| 60 | pred_value = pred_value.flatten() |
| 61 | new_priority = new_priority.flatten() |
| 62 | |
| 63 | summary_writer.add_scalar('{}_statistics/new_priority_mean'.format(tag), new_priority.mean(), step_count) |
| 64 | summary_writer.add_scalar('{}_statistics/new_priority_std'.format(tag), new_priority.std(), step_count) |
| 65 | |
| 66 | summary_writer.add_scalar('{}_statistics/target_value_prefix_mean'.format(tag), target_value_prefix.mean(), step_count) |
| 67 | summary_writer.add_scalar('{}_statistics/target_value_prefix_std'.format(tag), target_value_prefix.std(), step_count) |
| 68 | summary_writer.add_scalar('{}_statistics/pre_value_prefix_mean'.format(tag), pred_value_prefix.mean(), step_count) |