def optimizer_step(self, batch_info, device, model, rollout):
rollout = rollout.to_transitions()<EOL>observations = rollout.batch_tensor('<STR_LIT>')<EOL>returns = rollout.batch_tensor('<STR_LIT>')<EOL>policy_params = model.policy(observations)<EOL>policy_entropy = torch.mean(model.entropy(policy_params))<EOL>policy_loss = self.calc_policy_loss(model, policy_params, policy_entropy, rollout)<EOL>policy_grad = p2v(autograd.grad(policy_loss, model.policy_parameters(), retain_graph=True)).detach()<EOL>kl_divergence = torch.mean(model.kl_divergence(policy_params.detach(), policy_params))<EOL>kl_divergence_gradient = p2v(torch.autograd.grad(kl_divergence, model.policy_parameters(), create_graph=True))<EOL>step_direction = conjugate_gradient_method(<EOL>matrix_vector_operator=lambda x: self.fisher_vector_product(x, kl_divergence_gradient, model),<EOL>loss_gradient=-policy_grad,<EOL>nsteps=self.cg_iters<EOL>)<EOL>shs = <NUM_LIT:0.5> * step_direction @ self.fisher_vector_product(step_direction, kl_divergence_gradient, model)<EOL>lm = torch.sqrt(shs / self.mak_kl)<EOL>full_step = step_direction / lm<EOL>expected_improvement = (-policy_grad) @ full_step<EOL>original_parameter_vec = p2v(model.policy_parameters()).detach_()<EOL>policy_optimization_success, ratio, policy_loss_improvement, new_policy_loss, kl_divergence_step = self.line_search(<EOL>model, rollout, policy_loss, policy_params, original_parameter_vec, full_step, expected_improvement<EOL>)<EOL>gradient_norms = []<EOL>for i in range(self.vf_iters):<EOL><INDENT>batch_info.optimizer.zero_grad()<EOL>value_loss = self.value_loss(model, observations, returns)<EOL>value_loss.backward()<EOL>if self.max_grad_norm is not None:<EOL><INDENT>grad_norm = torch.nn.utils.clip_grad_norm_(<EOL>filter(lambda p: p.requires_grad, model.parameters()),<EOL>max_norm=self.max_grad_norm<EOL>)<EOL>gradient_norms.append(grad_norm)<EOL><DEDENT>batch_info.optimizer.step(closure=None)<EOL><DEDENT>if gradient_norms:<EOL><INDENT>gradient_norm = np.mean(gradient_norms)<EOL><DEDENT>else:<EOL><INDENT>gradient_norm = <NUM_LIT:0.0><EOL><DEDENT>return {<EOL>'<STR_LIT>': new_policy_loss.item(),<EOL>'<STR_LIT>': policy_entropy.item(),<EOL>'<STR_LIT>': value_loss.item(),<EOL>'<STR_LIT>': float(policy_optimization_success),<EOL>'<STR_LIT>': ratio.item(),<EOL>'<STR_LIT>': kl_divergence_step.item(),<EOL>'<STR_LIT>': policy_loss_improvement.item(),<EOL>'<STR_LIT>': gradient_norm,<EOL>'<STR_LIT>': torch.norm(rollout.batch_tensor('<STR_LIT>')).item(),<EOL>'<STR_LIT>': explained_variance(returns, rollout.batch_tensor('<STR_LIT>'))<EOL>}<EOL>
Single optimization step for a model