Single optimization step for a model

def optimizer_step(self, batch_info, device, model, rollout):
        """ Single optimization step for a model """
        rollout = rollout.to_transitions()

        # This algorithm makes quote strong assumptions about how does the model look
        # so it does not make that much sense to switch to the evaluator interface
        # As it would be more of a problem than actual benefit

        observations = rollout.batch_tensor('observations')
        returns = rollout.batch_tensor('returns')

        # Evaluate model on the observations
        policy_params = model.policy(observations)
        policy_entropy = torch.mean(model.entropy(policy_params))

        policy_loss = self.calc_policy_loss(model, policy_params, policy_entropy, rollout)
        policy_grad = p2v(autograd.grad(policy_loss, model.policy_parameters(), retain_graph=True)).detach()

        # Calculate gradient of KL divergence of model with fixed version of itself
        # Value of kl_divergence will be 0, but what we need is the gradient, actually the 2nd derivarive
        kl_divergence = torch.mean(model.kl_divergence(policy_params.detach(), policy_params))
        kl_divergence_gradient = p2v(torch.autograd.grad(kl_divergence, model.policy_parameters(), create_graph=True))

        step_direction = conjugate_gradient_method(
            matrix_vector_operator=lambda x: self.fisher_vector_product(x, kl_divergence_gradient, model),
            # Because we want to decrease the loss, we want to go into the direction of -gradient
            loss_gradient=-policy_grad,
            nsteps=self.cg_iters
        )

        shs = 0.5 * step_direction @ self.fisher_vector_product(step_direction, kl_divergence_gradient, model)
        lm = torch.sqrt(shs / self.mak_kl)
        full_step = step_direction / lm

        # Because we want to decrease the loss, we want to go into the direction of -gradient
        expected_improvement = (-policy_grad) @ full_step
        original_parameter_vec = p2v(model.policy_parameters()).detach_()

        policy_optimization_success, ratio, policy_loss_improvement, new_policy_loss, kl_divergence_step = self.line_search(
            model, rollout, policy_loss, policy_params, original_parameter_vec, full_step, expected_improvement
        )

        gradient_norms = []

        for i in range(self.vf_iters):
            batch_info.optimizer.zero_grad()
            value_loss = self.value_loss(model, observations, returns)

            value_loss.backward()

            # Gradient clipping
            if self.max_grad_norm is not None:
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    filter(lambda p: p.requires_grad, model.parameters()),
                    max_norm=self.max_grad_norm
                )

                gradient_norms.append(grad_norm)

            batch_info.optimizer.step(closure=None)

        if gradient_norms:
            gradient_norm = np.mean(gradient_norms)
        else:
            gradient_norm = 0.0

        # noinspection PyUnboundLocalVariable
        return {
            'new_policy_loss': new_policy_loss.item(),
            'policy_entropy': policy_entropy.item(),
            'value_loss': value_loss.item(),
            'policy_optimization_success': float(policy_optimization_success),
            'policy_improvement_ratio': ratio.item(),
            'kl_divergence_step': kl_divergence_step.item(),
            'policy_loss_improvement': policy_loss_improvement.item(),
            'grad_norm': gradient_norm,
            'advantage_norm': torch.norm(rollout.batch_tensor('advantages')).item(),
            'explained_variance': explained_variance(returns, rollout.batch_tensor('values'))
        }