def sample_from_rollouts(rollouts, baseline=None, reward_type="exact_match"):
  """Sample a single example from the given rollouts.

  Args:
    rollouts: Instance of RolloutOutputs.
    baseline: <float32> [batch_size] Baseline value b for R'(y) = R(y) - b.
    reward_type: Choice between indicator, exact_match, and F1.

  Returns:
    rollout: Instance of text_utils.TextInputs.
    reward: <float32> [batch_size]
  """
  batch_size = tensor_utils.shape(rollouts.token_ids, 0)
  rollout_length = tensor_utils.shape(rollouts.token_ids, 2)

  # Self-critical baseline.
  if baseline is None:
    baseline = tf.zeros([batch_size])

  # [batch_size, num_rollouts]
  rewards = rollouts.rewards[reward_type] - tf.expand_dims(baseline, 1)

  # Mask zero reward samples.
  masked_scores = tf.where(
      tf.not_equal(rewards, 0), tf.zeros_like(rollouts.scores),
      tf.ones_like(rollouts.scores) * -1e8)

  # [batch_size, 1]
  sample_idx = tf.distributions.Categorical(logits=masked_scores).sample()
  sample_idx = tf.reshape(sample_idx, [batch_size, 1])

  # [batch_size]
  reward = tf.reshape(tensor_utils.gather(rewards, sample_idx), [-1])

  # [batch_size, rollout_length]
  token_ids = tf.reshape(
      tensor_utils.gather(rollouts.token_ids, sample_idx), [batch_size, -1])
  mask = tf.reshape(
      tensor_utils.gather(rollouts.mask, sample_idx), [batch_size, -1])
  segment_ids = tf.zeros_like(token_ids)
  positions = tf.tile(
      tf.expand_dims(tf.range(rollout_length), 0), [batch_size, 1])

  # Create text input.
  rollout = text_utils.TextInputs(
      token_ids=token_ids,
      mask=mask,
      segment_ids=segment_ids,
      positions=positions)

  return rollout, reward
Sample a single example from the given rollouts.

  Args:
    rollouts: Instance of RolloutOutputs.
    baseline: <float32> [batch_size] Baseline value b for R'(y) = R(y) - b.
    reward_type: Choice between indicator, exact_match, and F1.

  Returns:
    rollout: Instance of text_utils.TextInputs.
    reward: <float32> [batch_size]
  
language/capwap/utils/reward_utils.py