import tensorflow as tf

    orig_indices = tf.range(
        self._sample_batch_size, dtype=relabel_indices.dtype)
    with tf.name_scope("relabelling"):
      # How often are the originally commanded goals most optimal?
      opt_indices = tf.argmax(logits_vec, axis=1)
      orig_is_opt = opt_indices == orig_indices
      orig_opt_frac = tf.reduce_mean(tf.cast(orig_is_opt, tf.float32))
      tf.compat.v2.summary.scalar(
          name="orig_task_optimal", data=orig_opt_frac, step=global_step)

      # How often is the relabelled goal optimal?
      # The relabel_indices are [B, 1], so we need to remove the extra dim.
