Calculate loss of the supplied rollout

def calculate_gradient(self, batch_info, device, model, rollout):
        """ Calculate loss of the supplied rollout """
        assert isinstance(rollout, Trajectories), "ACER algorithm requires trajectory input"

        local_epsilon = 1e-6

        evaluator = model.evaluate(rollout)

        actions = evaluator.get('rollout:actions')
        rollout_probabilities = torch.exp(evaluator.get('rollout:logprobs'))

        # We calculate the trust-region update with respect to the average model
        if self.trust_region:
            self.update_average_model(model)

        logprobs = evaluator.get('model:logprobs')
        q = evaluator.get('model:q')

        # Selected action values
        action_logprobs = select_indices(logprobs, actions)
        action_q = select_indices(q, actions)

        # We only want to propagate gradients through specific variables
        with torch.no_grad():
            model_probabilities = torch.exp(logprobs)

            # Importance sampling correction - we must find the quotient of probabilities
            rho = model_probabilities / (rollout_probabilities + local_epsilon)

            # Probability quotient only for selected actions
            actions_rho = select_indices(rho, actions)

            # Calculate policy state values
            model_state_values = (model_probabilities * q).sum(dim=1)

            trajectory_rewards = rollout.transition_tensors['rewards']
            trajectory_dones = rollout.transition_tensors['dones']

            q_retraced = self.retrace(
                trajectory_rewards,
                trajectory_dones,
                action_q.reshape(trajectory_rewards.size()),
                model_state_values.reshape(trajectory_rewards.size()),
                actions_rho.reshape(trajectory_rewards.size()),
                rollout.rollout_tensors['final_values']
            ).flatten()

            advantages = q_retraced - model_state_values
            importance_sampling_coefficient = torch.min(actions_rho, self.rho_cap * torch.ones_like(actions_rho))

            explained_variance = 1 - torch.var(q_retraced - action_q) / torch.var(q_retraced)

        # Entropy of the policy distribution
        policy_entropy = torch.mean(model.entropy(logprobs))
        policy_gradient_loss = -torch.mean(advantages * importance_sampling_coefficient * action_logprobs)

        # Policy gradient bias correction
        with torch.no_grad():
            advantages_bias_correction = q - model_state_values.view(model_probabilities.size(0), 1)
            bias_correction_coefficient = F.relu(1.0 - self.rho_cap / (rho + local_epsilon))

        # This sum is an expectation with respect to action probabilities according to model policy
        policy_gradient_bias_correction_gain = torch.sum(
            logprobs * bias_correction_coefficient * advantages_bias_correction * model_probabilities,
            dim=1
        )

        policy_gradient_bias_correction_loss = - torch.mean(policy_gradient_bias_correction_gain)

        policy_loss = policy_gradient_loss + policy_gradient_bias_correction_loss

        q_function_loss = 0.5 * F.mse_loss(action_q, q_retraced)

        if self.trust_region:
            with torch.no_grad():
                average_evaluator = self.average_model.evaluate(rollout)
                average_action_logits = average_evaluator.get('model:logprobs')

            actor_loss = policy_loss - self.entropy_coefficient * policy_entropy
            q_loss = self.q_coefficient * q_function_loss

            actor_gradient = torch.autograd.grad(-actor_loss, logprobs, retain_graph=True)[0]

            # kl_divergence = model.kl_divergence(average_action_logits, action_logits).mean()
            # kl_divergence_grad = torch.autograd.grad(kl_divergence, action_logits, retain_graph=True)

            # Analytically calculated derivative of KL divergence on logits
            # That makes it hardcoded for discrete action spaces
            kl_divergence_grad_symbolic = - torch.exp(average_action_logits) / logprobs.size(0)

            k_dot_g = (actor_gradient * kl_divergence_grad_symbolic).sum(dim=-1)
            k_dot_k = (kl_divergence_grad_symbolic ** 2).sum(dim=-1)

            adjustment = (k_dot_g - self.trust_region_delta) / k_dot_k
            adjustment_clipped = adjustment.clamp(min=0.0)

            actor_gradient_updated = actor_gradient - adjustment_clipped.view(adjustment_clipped.size(0), 1)

            # Populate gradient from the newly updated fn
            logprobs.backward(gradient=-actor_gradient_updated, retain_graph=True)
            q_loss.backward(retain_graph=True)
        else:
            # Just populate gradient from the loss
            loss = policy_loss + self.q_coefficient * q_function_loss - self.entropy_coefficient * policy_entropy

            loss.backward()

        return {
            'policy_loss': policy_loss.item(),
            'policy_gradient_loss': policy_gradient_loss.item(),
            'policy_gradient_bias_correction': policy_gradient_bias_correction_loss.item(),
            'avg_q_selected': action_q.mean().item(),
            'avg_q_retraced': q_retraced.mean().item(),
            'q_loss': q_function_loss.item(),
            'policy_entropy': policy_entropy.item(),
            'advantage_norm': torch.norm(advantages).item(),
            'explained_variance': explained_variance.item(),
            'model_prob_std': model_probabilities.std().item(),
            'rollout_prob_std': rollout_probabilities.std().item()
        }