Policy and Value optimizer step.

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):
  """Policy and Value optimizer step."""
  # 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)