import tensorflow as tf

      # 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))
        orig_q_vals = tf.gather_nd(logits_vec, indices)
        tf.compat.v2.summary.scalar(
            name="orig_q_vals",
            data=tf.reduce_mean(orig_q_vals),
            step=global_step,
        )

      # What are the average Q values of the relabelled tasks?
      indices = tf.transpose(
