code: def value_loss_given_predictions(value_prediction,
                                   rewards,
                                   reward_mask,
                                   gamma=0.99):
    
  
    B, T = rewards.shape  # pylint: disable=invalid-name
    assert (B, T) == reward_mask.shape
    assert (B, T + 1, 1) == value_prediction.shape
  
    value_prediction = np.squeeze(value_prediction, axis=2)  # (B, T+1)
    value_prediction = value_prediction[:, :-1] * reward_mask  # (B, T)
    r2g = rewards_to_go(rewards, reward_mask, gamma=gamma)  # (B, T)
    loss = (value_prediction - r2g)**2
  
    # Take an average on only the points where mask != 0.
    return np.sum(loss) / np.sum(reward_mask)
docs: Computes the value loss given the prediction of the value function.
  
    Args:
      value_prediction: np.ndarray of shape (B, T+1, 1)
      rewards: np.ndarray of shape (B, T) of rewards.
      reward_mask: np.ndarray of shape (B, T), the mask over rewards.
      gamma: float, discount factor.
  
    Returns:
      The average L2 value loss, averaged over instances where reward_mask is 1.