code: def ppo_loss(policy_net_apply,
               new_policy_params,
               old_policy_params,
               value_net_apply,
               value_net_params,
               padded_observations,
               padded_actions,
               padded_rewards,
               reward_mask,
               gamma=0.99,
               lambda_=0.95,
               epsilon=0.2):
    
    B, T = padded_rewards.shape  # pylint: disable=invalid-name
    assert (B, T + 1) == padded_observations.shape[:2]
    assert (B, T) == padded_actions.shape
    assert (B, T) == padded_rewards.shape
    assert (B, T) == reward_mask.shape
  
    # Compute predicted values and predicted log-probs and hand it over to
    # `ppo_loss_given_predictions`.
  
    # (B, T+1, 1)
    predicted_values = value_net_apply(padded_observations, value_net_params)
    assert (B, T + 1, 1) == predicted_values.shape
  
    # log_probab_actions_{old,new} are both (B, T+1, A)
    log_probab_actions_old = policy_net_apply(padded_observations,
                                              old_policy_params)
    log_probab_actions_new = policy_net_apply(padded_observations,
                                              new_policy_params)
    assert (B, T + 1) == log_probab_actions_old.shape[:2]
    assert (B, T + 1) == log_probab_actions_new.shape[:2]
    assert log_probab_actions_old.shape[-1] == log_probab_actions_new.shape[-1]
  
    return ppo_loss_given_predictions(log_probab_actions_new,
                                      log_probab_actions_old,
                                      predicted_values,
                                      padded_actions,
                                      padded_rewards,
                                      reward_mask,
                                      gamma=gamma,
                                      lambda_=lambda_,
                                      epsilon=epsilon)
docs: PPO objective, with an eventual minus sign, given observations.