(dist_inputs, values, returns, dones, rewards,
                 actions, old_log_probs, log_prob_fun, epsilon,
                 normalize_advantages)
return ppo_objective
PPO Objective.
PPO Objective.
PPO Objective .
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
def PPOObjective ( dist_inputs , values , returns , dones , rewards , actions , old_log_probs , log_prob_fun , epsilon , normalize_advantages ) : # 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