update models given a batch data Parameters ---------- model: Any EfficientZero models batch: Any a batch data inlcudes [inputs_batch, targets_batch] replay_buffer: Any replay buffer scaler: Any scaler for torch amp vis_result: bool
(model, batch, optimizer, replay_buffer, config, scaler, vis_result=False)
| 42 | |
| 43 | |
| 44 | def update_weights(model, batch, optimizer, replay_buffer, config, scaler, vis_result=False): |
| 45 | """update models given a batch data |
| 46 | Parameters |
| 47 | ---------- |
| 48 | model: Any |
| 49 | EfficientZero models |
| 50 | batch: Any |
| 51 | a batch data inlcudes [inputs_batch, targets_batch] |
| 52 | replay_buffer: Any |
| 53 | replay buffer |
| 54 | scaler: Any |
| 55 | scaler for torch amp |
| 56 | vis_result: bool |
| 57 | True -> log some visualization data in tensorboard (some distributions, values, etc) |
| 58 | """ |
| 59 | inputs_batch, targets_batch = batch |
| 60 | obs_batch_ori, action_batch, mask_batch, indices, weights_lst, make_time = inputs_batch |
| 61 | target_value_prefix, target_value, target_policy = targets_batch |
| 62 | |
| 63 | # [:, 0: config.stacked_observations * 3,:,:] |
| 64 | # obs_batch_ori is the original observations in a batch |
| 65 | # obs_batch is the observation for hat s_t (predicted hidden states from dynamics function) |
| 66 | # obs_target_batch is the observations for s_t (hidden states from representation function) |
| 67 | # to save GPU memory usage, obs_batch_ori contains (stack + unroll steps) frames |
| 68 | obs_batch_ori = torch.from_numpy(obs_batch_ori).to(config.device).float() / 255.0 |
| 69 | obs_batch = obs_batch_ori[:, 0: config.stacked_observations * config.image_channel, :, :] |
| 70 | obs_target_batch = obs_batch_ori[:, config.image_channel:, :, :] |
| 71 | |
| 72 | # do augmentations |
| 73 | if config.use_augmentation: |
| 74 | obs_batch = config.transform(obs_batch) |
| 75 | obs_target_batch = config.transform(obs_target_batch) |
| 76 | |
| 77 | # use GPU tensor |
| 78 | action_batch = torch.from_numpy(action_batch).to(config.device).unsqueeze(-1).long() |
| 79 | mask_batch = torch.from_numpy(mask_batch).to(config.device).float() |
| 80 | target_value_prefix = torch.from_numpy(target_value_prefix).to(config.device).float() |
| 81 | target_value = torch.from_numpy(target_value).to(config.device).float() |
| 82 | target_policy = torch.from_numpy(target_policy).to(config.device).float() |
| 83 | weights = torch.from_numpy(weights_lst).to(config.device).float() |
| 84 | |
| 85 | batch_size = obs_batch.size(0) |
| 86 | assert batch_size == config.batch_size == target_value_prefix.size(0) |
| 87 | metric_loss = torch.nn.L1Loss() |
| 88 | |
| 89 | # some logs preparation |
| 90 | other_log = {} |
| 91 | other_dist = {} |
| 92 | |
| 93 | other_loss = { |
| 94 | 'l1': -1, |
| 95 | 'l1_1': -1, |
| 96 | 'l1_-1': -1, |
| 97 | 'l1_0': -1, |
| 98 | } |
| 99 | for i in range(config.num_unroll_steps): |
| 100 | key = 'unroll_' + str(i + 1) + '_l1' |
| 101 | other_loss[key] = -1 |
no test coverage detected