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
Definition of the Advantage Actor Critic (A2C) loss.
trax/rl/rl_layers.py