import tensorflow as tf

          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.
      relabel_is_opt = tf.squeeze(relabel_indices) == orig_indices
      relabel_opt_frac = tf.reduce_mean(tf.cast(relabel_is_opt, tf.float32))
      tf.compat.v2.summary.scalar(
          name="relabel_task_optimal", data=relabel_opt_frac, step=global_step)

      # What are the average Q values of the original tasks?
      if batch_size == num_tasks:
        indices = tf.transpose(tf.stack([orig_indices, orig_indices], axis=0))
