code: def PPOObjective(
      dist_inputs,
      values,
      returns,
      dones,
      rewards,
      actions,
      old_log_probs,
      log_prob_fun,
      epsilon,
      normalize_advantages,
  ):
      """PPO Objective."""
      # dist_inputs of the shape float32[128,1,18]
      # values of the shape float32[128,1,1]
      # returns of the shape float32[128,1,1]
      # dones of the shape float32[128,1,1]
      # rewards of the shape int32[128,1,1]
      # actions of the shape int32[128,1]
      # and old_log_probs of the shape float32[128,1]
      returns = returns.squeeze(axis=2)
      values = values.squeeze(axis=2)
      dones = dones.squeeze(axis=2)
      rewards = rewards.squeeze(axis=2)
      assert (
          rewards.shape == dones.shape
      ), f"rewards.shape was {rewards.shape} and dones.shape was {dones.shape}"
      assert (
          dones.shape == values.shape
      ), f"dones.shape was {dones.shape} and values.shape was {values.shape}"
      assert (
          returns.shape == values.shape
      ), f"returns.shape was {returns.shape} and values.shape was {values.shape}"
      assert returns.shape == old_log_probs.shape, (
          f"returns.shape was {returns.shape} and"
          f"old_log_probs.shape was {old_log_probs.shape}"
      )
  
      probs_ratio = ProbsRatio(
          dist_inputs, actions, old_log_probs, log_prob_fun
      )
      assert probs_ratio.shape == old_log_probs.shape, (
          f"probs_ratio.shape was {probs_ratio.shape} and"
          f"old_log_probs.shape was {old_log_probs.shape}"
      )
  
      # jaxified versions of
      # returns[dones] = rewards[dones]
      # values[dones] = 0
      returns = jnp.where(dones, rewards, returns)
      values = jnp.where(dones, jnp.zeros_like(values), values)
      advantages = returns - values
      if normalize_advantages:
          advantages = advantages - jnp.mean(advantages)
          advantages /= jnp.std(advantages) + 1e-8
      assert old_log_probs.shape == advantages.shape, (
          f"old_log_probs.shape was {old_log_probs.shape} and advantages.shape was "
          f"{advantages.shape}"
      )
  
      unclipped_objective = UnclippedObjective(probs_ratio, advantages)
      assert unclipped_objective.shape == advantages.shape, (
          f"old_log_probs.shape was {old_log_probs.shape} and"
          f"unclipped_objective.shape was {unclipped_objective.shape}"
      )
  
      clipped_objective = ClippedObjective(
          probs_ratio, advantages, epsilon
      )
      assert clipped_objective.shape == advantages.shape, (
          f"clipped_objective.shape was {clipped_objective.shape} and"
          f"advantages.shape was {advantages.shape}"
      )
  
      ppo_objective = jnp.minimum(
          unclipped_objective, clipped_objective
      )
      assert ppo_objective.shape == advantages.shape, (
          f"ppo_objective.shape was {ppo_objective.shape} and"
          f"advantages.shape was {advantages.shape}"
      )
  
      return ppo_objective
  
file_path: trax/trax/rl/rl_layers.py
imports: from trax.fastmath import numpy
label_desc: Use image node_id 10 for calling a global function with example usage: PPOObjective(dist_inputs, values, returns, dones, rewards, actions, old_log_probs, log_prob_fun, epsilon, normalize_advantages) and returns: ppo_objective