def compute_rollouts(
    model,
    rc_model,
    features,
    encoder_cache,
    encoder_cache_mask,
    vocab,
    params,
):
  """Rollout model and compute rewards for each sample.

  Args:
    model: utils.transformer_utils.TransformerModel instance.
    rc_model: TF Hub module for extractive QA.
    features: Input features (questions and answers).
    encoder_cache: Transformer cache for encoded input.
    encoder_cache_mask: Input mask for the Transformer cache.
    vocab: Instance of text_utils.Vocab.
    params: Model parameters.

  Returns:
    rollout: Instance of RolloutOutputs.
  """
  # 1) First rollout the model with top-K beam search.
  rollout = transformer_utils.beam_search_decode(
      model=model,
      encoder_cache=encoder_cache,
      encoder_cache_mask=encoder_cache_mask,
      start_id=vocab.t2i(vocab.CLS),
      stop_id=vocab.t2i(vocab.SEP),
      segment_id=0,
      num_steps=params["decode_length"],
      beam_size=params["num_rollouts"],
      alpha=params["beam_length_penalty"],
      reuse=tf.AUTO_REUSE)

  # [batch_size, num_rollouts, rollout_length]
  batch_size = tensor_utils.shape(rollout.token_ids, 0)
  num_rollouts = tensor_utils.shape(rollout.token_ids, 1)
  rollout_ids = rollout.token_ids
  rollout_mask = rollout.mask

  # [batch_size * num_rollouts, rollout_length]
  rollout_length = tensor_utils.shape(rollout_ids, -1)
  rollout_ids = tf.reshape(rollout_ids, [-1, rollout_length])
  rollout_mask = tf.reshape(rollout_mask, [-1, rollout_length])

  # 2) Compute the QA rewards on the rollouts.
  # [batch_size * num_rollouts, question_length]
  question = tensor_utils.tile_batch(features["question_inputs"], num_rollouts)

  # [batch_size * num_rollouts, answer_length]
  answer = tensor_utils.tile_batch(features["answer_outputs"], num_rollouts)

  # [batch_size * num_rollouts]
  rewards = compute_qa_rewards(
      question_ids=question.token_ids,
      question_mask=question.mask,
      answer_ids=answer.token_ids,
      answer_mask=answer.mask,
      context_ids=rollout_ids[:, 1:],
      context_mask=rollout_mask[:, 1:],
      rc_model=rc_model,
      vocab=vocab,
      max_answer_length=params["answer_length"],
      no_answer_bias=params["no_answer_bias"])

  # [batch_size, num_rollouts, ...]
  reshaped_rewards = {}
  for k, v in rewards.items():
    if len(v.shape) > 1:
      v = tf.reshape(v, [batch_size, num_rollouts, -1])
    else:
      v = tf.reshape(v, [batch_size, num_rollouts])
    reshaped_rewards[k] = v

  # 3) Combine rollouts and rewards.
  rollouts = RolloutOutputs(
      token_ids=rollout.token_ids,
      mask=rollout.mask,
      scores=rollout.scores,
      rewards=reshaped_rewards)

  return rollouts
Rollout model and compute rewards for each sample.

  Args:
    model: utils.transformer_utils.TransformerModel instance.
    rc_model: TF Hub module for extractive QA.
    features: Input features (questions and answers).
    encoder_cache: Transformer cache for encoded input.
    encoder_cache_mask: Input mask for the Transformer cache.
    vocab: Instance of text_utils.Vocab.
    params: Model parameters.

  Returns:
    rollout: Instance of RolloutOutputs.
  
language/capwap/utils/reward_utils.py