code: def policy_and_value_opt_step(i,
                                opt_state,
                                opt_update,
                                policy_and_value_net_apply,
                                old_params,
                                padded_observations,
                                padded_actions,
                                padded_rewards,
                                reward_mask,
                                c1=1.0,
                                c2=0.01,
                                gamma=0.99,
                                lambda_=0.95,
                                epsilon=0.1):
    
    # Combined loss function given the new params.
    def policy_and_value_loss(params):
      """Returns the combined loss given just parameters."""
      (loss, _, _, _) = combined_loss(
          params,
          old_params,
          policy_and_value_net_apply,
          padded_observations,
          padded_actions,
          padded_rewards,
          reward_mask,
          c1=c1,
          c2=c2,
          gamma=gamma,
          lambda_=lambda_,
          epsilon=epsilon)
      return loss
  
    new_params = trax_opt.get_params(opt_state)
    g = grad(policy_and_value_loss)(new_params)
    return opt_update(i, g, opt_state)
docs: Policy and Value optimizer step.