Find the right stepsize to make sure policy improves

def line_search(self, model, rollout, original_policy_loss, original_policy_params, original_parameter_vec,
                    full_step, expected_improvement_full):
        """ Find the right stepsize to make sure policy improves """
        current_parameter_vec = original_parameter_vec.clone()

        for idx in range(self.line_search_iters):
            stepsize = 0.5 ** idx

            new_parameter_vec = current_parameter_vec + stepsize * full_step

            # Update model parameters
            v2p(new_parameter_vec, model.policy_parameters())

            # Calculate new loss
            with torch.no_grad():
                policy_params = model.policy(rollout.batch_tensor('observations'))
                policy_entropy = torch.mean(model.entropy(policy_params))
                kl_divergence = torch.mean(model.kl_divergence(original_policy_params, policy_params))

                new_loss = self.calc_policy_loss(model, policy_params, policy_entropy, rollout)

                actual_improvement = original_policy_loss - new_loss
                expected_improvement = expected_improvement_full * stepsize

                ratio = actual_improvement / expected_improvement

            if kl_divergence.item() > self.mak_kl * 1.5:
                # KL divergence bound exceeded
                continue
            elif ratio < expected_improvement:
                # Not enough loss improvement
                continue
            else:
                # Optimization successful
                return True, ratio, actual_improvement, new_loss, kl_divergence

        # Optimization failed, revert to initial parameters
        v2p(original_parameter_vec, model.policy_parameters())
        return False, torch.tensor(0.0), torch.tensor(0.0), torch.tensor(0.0), torch.tensor(0.0)