(dist_inputs, values, returns, dones, rewards,
                 actions, mask, log_prob_fun, normalize_advantages)
return a2c_objective
Definition of the Advantage Actor Critic (A2C) loss.
Definition of the Advantage Actor Critic (A2C) loss.
Definition of the Advantage Actor Critic ( A2C ) loss .
def A2CObjective(dist_inputs, values, returns, dones, rewards,
                 actions, mask, log_prob_fun, normalize_advantages):
  """Definition of the Advantage Actor Critic (A2C) loss."""
  # 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 int32[128,1,1]
  # actions of the shape int32[128,1]
  # and mask of the shape float32[128,1]
  # We have to squeeze values and returns, because we
  # are planning to compute (return - values) * new_log_probs * mask
  # and all of them should be of the same dimension
  values = values.squeeze(axis=2)
  returns = returns.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 values.shape == mask.shape, (
      f'values.shape was {values.shape} and mask.shape was {mask.shape}')
  assert returns.shape[0] == dist_inputs.shape[0], (
      f'returns.shape[0] was {returns.shape[0]} and dist_inputs.shape[0] was '
      f'{dist_inputs.shape[0]}')

  new_log_probs = NewLogProbs(dist_inputs, actions, log_prob_fun)
  assert new_log_probs.shape == mask.shape, (
      f'new_log_probs.shape was {new_log_probs.shape} and mask.shape was '
      f'{mask.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 new_log_probs.shape == advantages.shape, (
      f'new_log_probs.shape was {new_log_probs.shape} and advantages.shape was '
      f'{advantages.shape}')

  # One of the motivation to the squeezes and assertions is to
  # avoid [128,1] * [128,1,1] * [128] multiplications in the definition
  # of the a2c objective - we insist on the same shapes
  a2c_objective = -jnp.sum(new_log_probs * advantages * mask) / jnp.sum(mask)
  return a2c_objective
def A2CObjective ( dist_inputs , values , returns , dones , rewards , actions , mask , log_prob_fun , 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 int32[128,1,1] # actions of the shape int32[128,1] # and mask of the shape float32[128,1] # We have to squeeze values and returns, because we # are planning to compute (return - values) * new_log_probs * mask # and all of them should be of the same dimension values = values . squeeze ( axis = 2 ) returns = returns . 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 values . shape == mask . shape , ( f'values.shape was {values.shape} and mask.shape was {mask.shape}' ) assert returns . shape [ 0 ] == dist_inputs . shape [ 0 ] , ( f'returns.shape[0] was {returns.shape[0]} and dist_inputs.shape[0] was ' f'{dist_inputs.shape[0]}' ) new_log_probs = NewLogProbs ( dist_inputs , actions , log_prob_fun ) assert new_log_probs . shape == mask . shape , ( f'new_log_probs.shape was {new_log_probs.shape} and mask.shape was ' f'{mask.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 new_log_probs . shape == advantages . shape , ( f'new_log_probs.shape was {new_log_probs.shape} and advantages.shape was ' f'{advantages.shape}' ) # One of the motivation to the squeezes and assertions is to # avoid [128,1] * [128,1,1] * [128] multiplications in the definition # of the a2c objective - we insist on the same shapes a2c_objective = - jnp . sum ( new_log_probs * advantages * mask ) / jnp . sum ( mask ) return a2c_objective