def __call__( self, model: BaseNetwork, target_model: BaseNetwork, experiences: Tuple[torch.Tensor, ...], gamma: float, head_cfg: ConfigDict, ) -> Tuple[torch.Tensor, torch.Tensor]: """Return element-wise C51 loss and Q-values.""" states, actions, rewards, next_states, dones = experiences[:5] batch_size = states.shape[0] support = torch.linspace( head_cfg.configs.v_min, head_cfg.configs.v_max, head_cfg.configs.atom_size ).to(device) delta_z = float(head_cfg.configs.v_max - head_cfg.configs.v_min) / ( head_cfg.configs.atom_size - 1 ) with torch.no_grad(): # According to noisynet paper, # it resamples noisynet parameters on online network when using double q # but we don't because there is no remarkable difference in performance. next_actions = model.forward_(next_states)[1].argmax(1) next_dist = target_model.forward_(next_states)[0] next_dist = next_dist[range(batch_size), next_actions] t_z = rewards + (1 - dones) * gamma * support t_z = t_z.clamp(min=head_cfg.configs.v_min, max=head_cfg.configs.v_max) b = (t_z - head_cfg.configs.v_min) / delta_z l = b.floor().long() # noqa: E741 u = b.ceil().long() offset = ( torch.linspace( 0, (batch_size - 1) * head_cfg.configs.atom_size, batch_size ) .long() .unsqueeze(1) .expand(batch_size, head_cfg.configs.atom_size) .to(device) ) proj_dist = torch.zeros(next_dist.size(), device=device) proj_dist.view(-1).index_add_( 0, (l + offset).view(-1), (next_dist * (u.float() - b)).view(-1) ) proj_dist.view(-1).index_add_( 0, (u + offset).view(-1), (next_dist * (b - l.float())).view(-1) ) dist, q_values = model.forward_(states) log_p = torch.log(dist[range(batch_size), actions.long()]) dq_loss_element_wise = -(proj_dist * log_p).sum(1) return dq_loss_element_wise, q_values
def calculate_iqn_loss( model: BaseNetwork, target_model: BaseNetwork, experiences: Tuple[torch.Tensor, ...], gamma: float, batch_size: int, n_tau_samples: int, n_tau_prime_samples: int, kappa: float, ) -> Tuple[torch.Tensor, torch.Tensor]: """Return element-wise IQN loss and Q-values. Reference: https://github.com/google/dopamine """ states, actions, rewards, next_states, dones = experiences[:5] # size of rewards: (n_tau_prime_samples x batch_size) x 1. rewards = rewards.repeat(n_tau_prime_samples, 1) # size of gamma_with_terminal: (n_tau_prime_samples x batch_size) x 1. masks = 1 - dones gamma_with_terminal = masks * gamma gamma_with_terminal = gamma_with_terminal.repeat(n_tau_prime_samples, 1) # Get the indices of the maximium Q-value across the action dimension. # Shape of replay_next_qt_argmax: (n_tau_prime_samples x batch_size) x 1. next_actions = model(next_states).argmax(dim=1) # double Q next_actions = next_actions[:, None] next_actions = next_actions.repeat(n_tau_prime_samples, 1) # Shape of next_target_values: (n_tau_prime_samples x batch_size) x 1. target_quantile_values, _ = target_model.forward_(next_states, n_tau_prime_samples) target_quantile_values = target_quantile_values.gather(1, next_actions) target_quantile_values = rewards + gamma_with_terminal * target_quantile_values target_quantile_values = target_quantile_values.detach() # Reshape to n_tau_prime_samples x batch_size x 1 since this is # the manner in which the target_quantile_values are tiled. target_quantile_values = target_quantile_values.view( n_tau_prime_samples, batch_size, 1) # Transpose dimensions so that the dimensionality is batch_size x # n_tau_prime_samples x 1 to prepare for computation of Bellman errors. target_quantile_values = torch.transpose(target_quantile_values, 0, 1) # Get quantile values: (n_tau_samples x batch_size) x action_dim. quantile_values, quantiles = model.forward_(states, n_tau_samples) reshaped_actions = actions[:, None].repeat(n_tau_samples, 1) chosen_action_quantile_values = quantile_values.gather( 1, reshaped_actions.long()) chosen_action_quantile_values = chosen_action_quantile_values.view( n_tau_samples, batch_size, 1) # Transpose dimensions so that the dimensionality is batch_size x # n_tau_prime_samples x 1 to prepare for computation of Bellman errors. chosen_action_quantile_values = torch.transpose( chosen_action_quantile_values, 0, 1) # Shape of bellman_erors and huber_loss: # batch_size x num_tau_prime_samples x num_tau_samples x 1. bellman_errors = (target_quantile_values[:, :, None, :] - chosen_action_quantile_values[:, None, :, :]) # The huber loss (introduced in QR-DQN) is defined via two cases: # case_one: |bellman_errors| <= kappa # case_two: |bellman_errors| > kappa huber_loss_case_one = ((torch.abs(bellman_errors) <= kappa).float() * 0.5 * bellman_errors**2) huber_loss_case_two = ((torch.abs(bellman_errors) > kappa).float() * kappa * (torch.abs(bellman_errors) - 0.5 * kappa)) huber_loss = huber_loss_case_one + huber_loss_case_two # Reshape quantiles to batch_size x num_tau_samples x 1 quantiles = quantiles.view(n_tau_samples, batch_size, 1) quantiles = torch.transpose(quantiles, 0, 1) # Tile by num_tau_prime_samples along a new dimension. Shape is now # batch_size x num_tau_prime_samples x num_tau_samples x 1. # These quantiles will be used for computation of the quantile huber loss # below (see section 2.3 of the paper). quantiles = quantiles[:, None, :, :].repeat(1, n_tau_prime_samples, 1, 1) quantiles = quantiles.to(device) # Shape: batch_size x n_tau_prime_samples x n_tau_samples x 1. quantile_huber_loss = (torch.abs(quantiles - (bellman_errors < 0).float().detach()) * huber_loss / kappa) # Sum over current quantile value (n_tau_samples) dimension, # average over target quantile value (n_tau_prime_samples) dimension. # Shape: batch_size x n_tau_prime_samples x 1. loss = torch.sum(quantile_huber_loss, dim=2) # Shape: batch_size x 1. iqn_loss_element_wise = torch.mean(loss, dim=1) # q values for regularization. q_values = model(states) return iqn_loss_element_wise, q_values