code: def value_loss(value_net_apply,
                 value_net_params,
                 observations,
                 rewards,
                 reward_mask,
                 gamma=0.99):
    
  
    B, T = rewards.shape  # pylint: disable=invalid-name
    assert (B, T + 1) == observations.shape[:2]
  
    # NOTE: observations is (B, T+1) + OBS, value_prediction is (B, T+1, 1)
    value_prediction = value_net_apply(observations, value_net_params)
    assert (B, T + 1, 1) == value_prediction.shape
  
    return value_loss_given_predictions(value_prediction, rewards, reward_mask,
                                        gamma)
docs: Computes the value loss.
  
    Args:
      value_net_apply: value net apply function with signature (params, ndarray of
        shape (B, T+1) + OBS) -> ndarray(B, T+1, 1)
      value_net_params: params of value_net_apply.
      observations: np.ndarray of shape (B, T+1) + OBS
      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.