def calculate_gradient(self, batch_info, device, model, rollout):
evaluator = model.evaluate(rollout)<EOL>batch_size = rollout.frames()<EOL>dones_tensor = evaluator.get('<STR_LIT>')<EOL>rewards_tensor = evaluator.get('<STR_LIT>')<EOL>assert dones_tensor.dtype == torch.float32<EOL>with torch.no_grad():<EOL><INDENT>target_evaluator = self.target_model.evaluate(rollout)<EOL>if self.double_dqn:<EOL><INDENT>target_value_histogram_for_all_actions = target_evaluator.get('<STR_LIT>').exp()<EOL>model_value_histogram_for_all_actions = evaluator.get('<STR_LIT>').exp()<EOL>atoms_aligned = self.support_atoms.view(<NUM_LIT:1>, <NUM_LIT:1>, self.num_atoms)<EOL>selected_action_indices = (<EOL>(atoms_aligned * model_value_histogram_for_all_actions).sum(dim=-<NUM_LIT:1>).argmax(dim=<NUM_LIT:1>)<EOL>)<EOL>next_value_histograms = (<EOL>target_value_histogram_for_all_actions[range(batch_size), selected_action_indices]<EOL>)<EOL><DEDENT>else:<EOL><INDENT>target_value_histogram_for_all_actions = target_evaluator.get('<STR_LIT>').exp()<EOL>atoms_aligned = self.support_atoms.view(<NUM_LIT:1>, <NUM_LIT:1>, self.num_atoms)<EOL>selected_action_indices = (<EOL>(atoms_aligned * target_value_histogram_for_all_actions).sum(dim=-<NUM_LIT:1>).argmax(dim=<NUM_LIT:1>)<EOL>)<EOL>next_value_histograms = (<EOL>target_value_histogram_for_all_actions[range(batch_size), selected_action_indices]<EOL>)<EOL><DEDENT>forward_steps = rollout.extra_data.get('<STR_LIT>', <NUM_LIT:1>)<EOL>atoms_projected = (<EOL>rewards_tensor.unsqueeze(<NUM_LIT:1>) +<EOL>(self.discount_factor ** forward_steps) *<EOL>(<NUM_LIT:1> - dones_tensor).unsqueeze(<NUM_LIT:1>) * self.support_atoms.unsqueeze(<NUM_LIT:0>)<EOL>)<EOL>atoms_projected = atoms_projected.clamp(min=self.vmin, max=self.vmax)<EOL>projection_indices = (atoms_projected - self.vmin) / self.atom_delta<EOL>index_floor = projection_indices.floor().long()<EOL>index_ceil = projection_indices.ceil().long()<EOL>index_floor[(index_ceil > <NUM_LIT:0>) * (index_floor == index_ceil)] -= <NUM_LIT:1><EOL>index_ceil[(index_floor < (self.num_atoms - <NUM_LIT:1>)) * (index_floor == index_ceil)] += <NUM_LIT:1><EOL>value_histogram_projected = torch.zeros_like(next_value_histograms)<EOL>offsets = (<EOL>torch.arange(<NUM_LIT:0>, batch_size * self.num_atoms, self.num_atoms)<EOL>.unsqueeze(<NUM_LIT:1>)<EOL>.expand(batch_size, self.num_atoms)<EOL>.contiguous().view(-<NUM_LIT:1>).to(device)<EOL>)<EOL>value_histogram_projected = value_histogram_projected.view(-<NUM_LIT:1>)<EOL>index_ceil = index_ceil.view(-<NUM_LIT:1>)<EOL>index_floor = index_floor.view(-<NUM_LIT:1>)<EOL>projection_indices = projection_indices.view(-<NUM_LIT:1>)<EOL>value_histogram_projected.index_add_(<EOL><NUM_LIT:0>,<EOL>index_floor+offsets,<EOL>(next_value_histograms.view(-<NUM_LIT:1>) * (index_ceil.float() - projection_indices))<EOL>)<EOL>value_histogram_projected.index_add_(<EOL><NUM_LIT:0>,<EOL>index_ceil+offsets,<EOL>(next_value_histograms.view(-<NUM_LIT:1>) * (projection_indices - index_floor.float()))<EOL>)<EOL>value_histogram_projected = value_histogram_projected.reshape(next_value_histograms.shape)<EOL><DEDENT>q_log_histogram_selected = evaluator.get('<STR_LIT>')<EOL>original_losses = -(value_histogram_projected * q_log_histogram_selected).sum(dim=<NUM_LIT:1>)<EOL>if evaluator.is_provided('<STR_LIT>'):<EOL><INDENT>weights = evaluator.get('<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>weights = torch.ones_like(rewards_tensor)<EOL><DEDENT>loss_value = torch.mean(weights * original_losses)<EOL>loss_value.backward()<EOL>with torch.no_grad():<EOL><INDENT>mean_q_model = (self.support_atoms.unsqueeze(<NUM_LIT:0>) * torch.exp(q_log_histogram_selected)).sum(dim=<NUM_LIT:1>).mean()<EOL>mean_q_target = (self.support_atoms.unsqueeze(<NUM_LIT:0>) * value_histogram_projected).sum(dim=<NUM_LIT:1>).mean()<EOL><DEDENT>return {<EOL>'<STR_LIT>': loss_value.item(),<EOL>'<STR_LIT>': original_losses.detach().cpu().numpy(),<EOL>'<STR_LIT>': mean_q_model.item(),<EOL>'<STR_LIT>': mean_q_target.item()<EOL>}<EOL>
Calculate loss of the supplied rollout